ominfra 0.0.0.dev137__py3-none-any.whl → 0.0.0.dev138__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
@@ -4,6 +4,10 @@
4
4
  # @omlish-script
5
5
  # @omlish-amalg-output main.py
6
6
  # ruff: noqa: N802 UP006 UP007 UP036
7
+ """
8
+ manage.py -s 'docker run -i python:3.12'
9
+ manage.py -qs 'ssh -i foo/bar foo@bar.baz' --python=python3.8
10
+ """
7
11
  import abc
8
12
  import base64
9
13
  import collections.abc
@@ -25,7 +29,6 @@ import site
25
29
  import struct
26
30
  import subprocess
27
31
  import sys
28
- import textwrap
29
32
  import threading
30
33
  import time
31
34
  import types
@@ -87,11 +90,120 @@ Basically this: https://mitogen.networkgenomics.com/howitworks.html
87
90
  ##
88
91
 
89
92
 
93
+ @dc.dataclass(frozen=True)
94
+ class PyremoteBootstrapOptions:
95
+ debug: bool = False
96
+
97
+
98
+ ##
99
+
100
+
101
+ @dc.dataclass(frozen=True)
102
+ class PyremoteEnvInfo:
103
+ sys_base_prefix: str
104
+ sys_byteorder: str
105
+ sys_defaultencoding: str
106
+ sys_exec_prefix: str
107
+ sys_executable: str
108
+ sys_implementation_name: str
109
+ sys_path: ta.List[str]
110
+ sys_platform: str
111
+ sys_prefix: str
112
+ sys_version: str
113
+ sys_version_info: ta.List[ta.Union[int, str]]
114
+
115
+ platform_architecture: ta.List[str]
116
+ platform_machine: str
117
+ platform_platform: str
118
+ platform_processor: str
119
+ platform_system: str
120
+ platform_release: str
121
+ platform_version: str
122
+
123
+ site_userbase: str
124
+
125
+ os_cwd: str
126
+ os_gid: int
127
+ os_loadavg: ta.List[float]
128
+ os_login: ta.Optional[str]
129
+ os_pgrp: int
130
+ os_pid: int
131
+ os_ppid: int
132
+ os_uid: int
133
+
134
+ pw_name: str
135
+ pw_uid: int
136
+ pw_gid: int
137
+ pw_gecos: str
138
+ pw_dir: str
139
+ pw_shell: str
140
+
141
+ env_path: ta.Optional[str]
142
+
143
+
144
+ def _get_pyremote_env_info() -> PyremoteEnvInfo:
145
+ os_uid = os.getuid()
146
+
147
+ pw = pwd.getpwuid(os_uid)
148
+
149
+ os_login: ta.Optional[str]
150
+ try:
151
+ os_login = os.getlogin()
152
+ except OSError:
153
+ os_login = None
154
+
155
+ return PyremoteEnvInfo(
156
+ sys_base_prefix=sys.base_prefix,
157
+ sys_byteorder=sys.byteorder,
158
+ sys_defaultencoding=sys.getdefaultencoding(),
159
+ sys_exec_prefix=sys.exec_prefix,
160
+ sys_executable=sys.executable,
161
+ sys_implementation_name=sys.implementation.name,
162
+ sys_path=sys.path,
163
+ sys_platform=sys.platform,
164
+ sys_prefix=sys.prefix,
165
+ sys_version=sys.version,
166
+ sys_version_info=list(sys.version_info),
167
+
168
+ platform_architecture=list(platform.architecture()),
169
+ platform_machine=platform.machine(),
170
+ platform_platform=platform.platform(),
171
+ platform_processor=platform.processor(),
172
+ platform_system=platform.system(),
173
+ platform_release=platform.release(),
174
+ platform_version=platform.version(),
175
+
176
+ site_userbase=site.getuserbase(),
177
+
178
+ os_cwd=os.getcwd(),
179
+ os_gid=os.getgid(),
180
+ os_loadavg=list(os.getloadavg()),
181
+ os_login=os_login,
182
+ os_pgrp=os.getpgrp(),
183
+ os_pid=os.getpid(),
184
+ os_ppid=os.getppid(),
185
+ os_uid=os_uid,
186
+
187
+ pw_name=pw.pw_name,
188
+ pw_uid=pw.pw_uid,
189
+ pw_gid=pw.pw_gid,
190
+ pw_gecos=pw.pw_gecos,
191
+ pw_dir=pw.pw_dir,
192
+ pw_shell=pw.pw_shell,
193
+
194
+ env_path=os.environ.get('PATH'),
195
+ )
196
+
197
+
198
+ ##
199
+
200
+
90
201
  _PYREMOTE_BOOTSTRAP_INPUT_FD = 100
91
202
  _PYREMOTE_BOOTSTRAP_SRC_FD = 101
92
203
 
93
204
  _PYREMOTE_BOOTSTRAP_CHILD_PID_VAR = '_OPYR_CHILD_PID'
94
205
  _PYREMOTE_BOOTSTRAP_ARGV0_VAR = '_OPYR_ARGV0'
206
+ _PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR = '_OPYR_OPTIONS_JSON'
95
207
 
96
208
  _PYREMOTE_BOOTSTRAP_ACK0 = b'OPYR000\n'
97
209
  _PYREMOTE_BOOTSTRAP_ACK1 = b'OPYR001\n'
@@ -153,7 +265,9 @@ def _pyremote_bootstrap_main(context_name: str) -> None:
153
265
 
154
266
  # Read main src from stdin
155
267
  main_z_len = struct.unpack('<I', os.read(0, 4))[0]
156
- main_src = zlib.decompress(os.fdopen(0, 'rb').read(main_z_len))
268
+ if len(main_z := os.fdopen(0, 'rb').read(main_z_len)) != main_z_len:
269
+ raise EOFError
270
+ main_src = zlib.decompress(main_z)
157
271
 
158
272
  # Write both copies of main src. Must write to w0 (parent stdin) before w1 (copy pipe) as pipe will likely fill
159
273
  # and block and need to be drained by pyremote_bootstrap_finalize running in parent.
@@ -176,6 +290,8 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
176
290
  if any(c in context_name for c in '\'"'):
177
291
  raise NameError(context_name)
178
292
 
293
+ import inspect
294
+ import textwrap
179
295
  bs_src = textwrap.dedent(inspect.getsource(_pyremote_bootstrap_main))
180
296
 
181
297
  for gl in [
@@ -202,9 +318,6 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
202
318
  bs_z = zlib.compress(bs_src.encode('utf-8'))
203
319
  bs_z64 = base64.encodebytes(bs_z).replace(b'\n', b'')
204
320
 
205
- def dq_repr(o: ta.Any) -> str:
206
- return '"' + repr(o)[1:-1] + '"'
207
-
208
321
  stmts = [
209
322
  f'import {", ".join(_PYREMOTE_BOOTSTRAP_IMPORTS)}',
210
323
  f'exec(zlib.decompress(base64.decodebytes({bs_z64!r})))',
@@ -219,99 +332,83 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
219
332
 
220
333
 
221
334
  @dc.dataclass(frozen=True)
222
- class PyremoteEnvInfo:
223
- sys_base_prefix: str
224
- sys_byteorder: str
225
- sys_defaultencoding: str
226
- sys_exec_prefix: str
227
- sys_executable: str
228
- sys_implementation_name: str
229
- sys_path: ta.List[str]
230
- sys_platform: str
231
- sys_prefix: str
232
- sys_version: str
233
- sys_version_info: ta.List[ta.Union[int, str]]
234
-
235
- platform_architecture: ta.List[str]
236
- platform_machine: str
237
- platform_platform: str
238
- platform_processor: str
239
- platform_system: str
240
- platform_release: str
241
- platform_version: str
242
-
243
- site_userbase: str
244
-
245
- os_cwd: str
246
- os_gid: int
247
- os_loadavg: ta.List[float]
248
- os_login: ta.Optional[str]
249
- os_pgrp: int
250
- os_pid: int
251
- os_ppid: int
252
- os_uid: int
253
-
254
- pw_name: str
255
- pw_uid: int
256
- pw_gid: int
257
- pw_gecos: str
258
- pw_dir: str
259
- pw_shell: str
260
-
261
- env_path: ta.Optional[str]
335
+ class PyremotePayloadRuntime:
336
+ input: ta.BinaryIO
337
+ output: ta.BinaryIO
338
+ main_src: str
339
+ options: PyremoteBootstrapOptions
340
+ env_info: PyremoteEnvInfo
262
341
 
263
342
 
264
- def _get_pyremote_env_info() -> PyremoteEnvInfo:
265
- os_uid = os.getuid()
343
+ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
344
+ # If json options var is not present we need to do initial finalization
345
+ if _PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR not in os.environ:
346
+ # Read second copy of main src
347
+ r1 = os.fdopen(_PYREMOTE_BOOTSTRAP_SRC_FD, 'rb', 0)
348
+ main_src = r1.read().decode('utf-8')
349
+ r1.close()
350
+
351
+ # Reap boostrap child. Must be done after reading second copy of source because source may be too big to fit in
352
+ # a pipe at once.
353
+ os.waitpid(int(os.environ.pop(_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR)), 0)
354
+
355
+ # Read options
356
+ options_json_len = struct.unpack('<I', os.read(_PYREMOTE_BOOTSTRAP_INPUT_FD, 4))[0]
357
+ if len(options_json := os.read(_PYREMOTE_BOOTSTRAP_INPUT_FD, options_json_len)) != options_json_len:
358
+ raise EOFError
359
+ options = PyremoteBootstrapOptions(**json.loads(options_json.decode('utf-8')))
360
+
361
+ # If debugging, re-exec as file
362
+ if options.debug:
363
+ # Write temp source file
364
+ import tempfile
365
+ tfd, tfn = tempfile.mkstemp('-pyremote.py')
366
+ os.write(tfd, main_src.encode('utf-8'))
367
+ os.close(tfd)
368
+
369
+ # Set json options var
370
+ os.environ[_PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR] = options_json.decode('utf-8')
371
+
372
+ # Re-exec temp file
373
+ os.execl(os.environ[_PYREMOTE_BOOTSTRAP_ARGV0_VAR], sys.orig_argv[0], tfn)
266
374
 
267
- pw = pwd.getpwuid(os_uid)
375
+ else:
376
+ # Load options json var
377
+ options_json_str = os.environ.pop(_PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR)
378
+ options = PyremoteBootstrapOptions(**json.loads(options_json_str))
268
379
 
269
- os_login: ta.Optional[str]
270
- try:
271
- os_login = os.getlogin()
272
- except OSError:
273
- os_login = None
380
+ # Read temp source file
381
+ with open(sys.orig_argv[1]) as sf:
382
+ main_src = sf.read()
274
383
 
275
- return PyremoteEnvInfo(
276
- sys_base_prefix=sys.base_prefix,
277
- sys_byteorder=sys.byteorder,
278
- sys_defaultencoding=sys.getdefaultencoding(),
279
- sys_exec_prefix=sys.exec_prefix,
280
- sys_executable=sys.executable,
281
- sys_implementation_name=sys.implementation.name,
282
- sys_path=sys.path,
283
- sys_platform=sys.platform,
284
- sys_prefix=sys.prefix,
285
- sys_version=sys.version,
286
- sys_version_info=list(sys.version_info),
384
+ # Restore original argv0
385
+ sys.executable = os.environ.pop(_PYREMOTE_BOOTSTRAP_ARGV0_VAR)
287
386
 
288
- platform_architecture=list(platform.architecture()),
289
- platform_machine=platform.machine(),
290
- platform_platform=platform.platform(),
291
- platform_processor=platform.processor(),
292
- platform_system=platform.system(),
293
- platform_release=platform.release(),
294
- platform_version=platform.version(),
387
+ # Write third ack
388
+ os.write(1, _PYREMOTE_BOOTSTRAP_ACK2)
295
389
 
296
- site_userbase=site.getuserbase(),
390
+ # Write env info
391
+ env_info = _get_pyremote_env_info()
392
+ env_info_json = json.dumps(dc.asdict(env_info), indent=None, separators=(',', ':')) # noqa
393
+ os.write(1, struct.pack('<I', len(env_info_json)))
394
+ os.write(1, env_info_json.encode('utf-8'))
297
395
 
298
- os_cwd=os.getcwd(),
299
- os_gid=os.getgid(),
300
- os_loadavg=list(os.getloadavg()),
301
- os_login=os_login,
302
- os_pgrp=os.getpgrp(),
303
- os_pid=os.getpid(),
304
- os_ppid=os.getppid(),
305
- os_uid=os_uid,
396
+ # Setup IO
397
+ input = os.fdopen(_PYREMOTE_BOOTSTRAP_INPUT_FD, 'rb', 0) # noqa
398
+ output = os.fdopen(os.dup(1), 'wb', 0) # noqa
399
+ os.dup2(nfd := os.open('/dev/null', os.O_WRONLY), 1)
400
+ os.close(nfd)
306
401
 
307
- pw_name=pw.pw_name,
308
- pw_uid=pw.pw_uid,
309
- pw_gid=pw.pw_gid,
310
- pw_gecos=pw.pw_gecos,
311
- pw_dir=pw.pw_dir,
312
- pw_shell=pw.pw_shell,
402
+ # Write fourth ack
403
+ output.write(_PYREMOTE_BOOTSTRAP_ACK3)
313
404
 
314
- env_path=os.environ.get('PATH'),
405
+ # Return
406
+ return PyremotePayloadRuntime(
407
+ input=input,
408
+ output=output,
409
+ main_src=main_src,
410
+ options=options,
411
+ env_info=env_info,
315
412
  )
316
413
 
317
414
 
@@ -319,12 +416,15 @@ def _get_pyremote_env_info() -> PyremoteEnvInfo:
319
416
 
320
417
 
321
418
  class PyremoteBootstrapDriver:
322
- def __init__(self, main_src: str) -> None:
419
+ def __init__(self, main_src: str, options: PyremoteBootstrapOptions = PyremoteBootstrapOptions()) -> None:
323
420
  super().__init__()
324
421
 
325
422
  self._main_src = main_src
326
423
  self._main_z = zlib.compress(main_src.encode('utf-8'))
327
424
 
425
+ self._options = options
426
+ self._options_json = json.dumps(dc.asdict(options), indent=None, separators=(',', ':')).encode('utf-8') # noqa
427
+
328
428
  #
329
429
 
330
430
  @dc.dataclass(frozen=True)
@@ -344,7 +444,7 @@ class PyremoteBootstrapDriver:
344
444
  env_info: PyremoteEnvInfo
345
445
 
346
446
  def gen(self) -> ta.Generator[ta.Union[Read, Write], ta.Optional[bytes], Result]:
347
- # Read first ack
447
+ # Read first ack (after fork)
348
448
  yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK0)
349
449
 
350
450
  # Read pid
@@ -355,8 +455,14 @@ class PyremoteBootstrapDriver:
355
455
  yield from self._write(struct.pack('<I', len(self._main_z)))
356
456
  yield from self._write(self._main_z)
357
457
 
358
- # Read second and third ack
458
+ # Read second ack (after writing src copies)
359
459
  yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK1)
460
+
461
+ # Write options
462
+ yield from self._write(struct.pack('<I', len(self._options_json)))
463
+ yield from self._write(self._options_json)
464
+
465
+ # Read third ack (after reaping child process)
360
466
  yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK2)
361
467
 
362
468
  # Read env info
@@ -366,7 +472,7 @@ class PyremoteBootstrapDriver:
366
472
  env_info_json = d.decode('utf-8')
367
473
  env_info = PyremoteEnvInfo(**json.loads(env_info_json))
368
474
 
369
- # Read fourth ack
475
+ # Read fourth ack (after finalization completed)
370
476
  yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK3)
371
477
 
372
478
  # Return
@@ -409,7 +515,8 @@ class PyremoteBootstrapDriver:
409
515
  return e.value
410
516
 
411
517
  if isinstance(go, self.Read):
412
- gi = stdout.read(go.sz)
518
+ if len(gi := stdout.read(go.sz)) != go.sz:
519
+ raise EOFError
413
520
  elif isinstance(go, self.Write):
414
521
  gi = None
415
522
  stdin.write(go.d)
@@ -418,57 +525,6 @@ class PyremoteBootstrapDriver:
418
525
  raise TypeError(go)
419
526
 
420
527
 
421
- ##
422
-
423
-
424
- @dc.dataclass(frozen=True)
425
- class PyremotePayloadRuntime:
426
- input: ta.BinaryIO
427
- output: ta.BinaryIO
428
- main_src: str
429
- env_info: PyremoteEnvInfo
430
-
431
-
432
- def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
433
- # Restore original argv0
434
- sys.executable = os.environ.pop(_PYREMOTE_BOOTSTRAP_ARGV0_VAR)
435
-
436
- # Read second copy of main src
437
- r1 = os.fdopen(_PYREMOTE_BOOTSTRAP_SRC_FD, 'rb', 0)
438
- main_src = r1.read().decode('utf-8')
439
- r1.close()
440
-
441
- # Reap boostrap child. Must be done after reading second copy of source because source may be too big to fit in a
442
- # pipe at once.
443
- os.waitpid(int(os.environ.pop(_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR)), 0)
444
-
445
- # Write third ack
446
- os.write(1, _PYREMOTE_BOOTSTRAP_ACK2)
447
-
448
- # Write env info
449
- env_info = _get_pyremote_env_info()
450
- env_info_json = json.dumps(dc.asdict(env_info), indent=None, separators=(',', ':')) # noqa
451
- os.write(1, struct.pack('<I', len(env_info_json)))
452
- os.write(1, env_info_json.encode('utf-8'))
453
-
454
- # Setup IO
455
- input = os.fdopen(_PYREMOTE_BOOTSTRAP_INPUT_FD, 'rb', 0) # noqa
456
- output = os.fdopen(os.dup(1), 'wb', 0) # noqa
457
- os.dup2(nfd := os.open('/dev/null', os.O_WRONLY), 1)
458
- os.close(nfd)
459
-
460
- # Write fourth ack
461
- output.write(_PYREMOTE_BOOTSTRAP_ACK3)
462
-
463
- # Return
464
- return PyremotePayloadRuntime(
465
- input=input,
466
- output=output,
467
- main_src=main_src,
468
- env_info=env_info,
469
- )
470
-
471
-
472
528
  ########################################
473
529
  # ../../../../omlish/lite/cached.py
474
530
 
@@ -1509,8 +1565,41 @@ class SubprocessCommand(Command['SubprocessCommand.Input', 'SubprocessCommand.Ou
1509
1565
  ##
1510
1566
 
1511
1567
 
1512
- def _send_obj(f: ta.IO, o: ta.Any) -> None:
1513
- j = json_dumps_compact(marshal_obj(o))
1568
+ _COMMAND_TYPES = {
1569
+ 'subprocess': SubprocessCommand,
1570
+ }
1571
+
1572
+
1573
+ register_opj_marshaler(
1574
+ Command.Input,
1575
+ PolymorphicObjMarshaler.of([
1576
+ PolymorphicObjMarshaler.Impl(
1577
+ cty.Input,
1578
+ k,
1579
+ get_obj_marshaler(cty.Input),
1580
+ )
1581
+ for k, cty in _COMMAND_TYPES.items()
1582
+ ]),
1583
+ )
1584
+
1585
+ register_opj_marshaler(
1586
+ Command.Output,
1587
+ PolymorphicObjMarshaler.of([
1588
+ PolymorphicObjMarshaler.Impl(
1589
+ cty.Output,
1590
+ k,
1591
+ get_obj_marshaler(cty.Output),
1592
+ )
1593
+ for k, cty in _COMMAND_TYPES.items()
1594
+ ]),
1595
+ )
1596
+
1597
+
1598
+ ##
1599
+
1600
+
1601
+ def _send_obj(f: ta.IO, o: ta.Any, ty: ta.Any = None) -> None:
1602
+ j = json_dumps_compact(marshal_obj(o, ty))
1514
1603
  d = j.encode('utf-8')
1515
1604
 
1516
1605
  f.write(struct.pack('<I', len(d)))
@@ -1518,17 +1607,17 @@ def _send_obj(f: ta.IO, o: ta.Any) -> None:
1518
1607
  f.flush()
1519
1608
 
1520
1609
 
1521
- def _recv_obj(f: ta.IO, ty: type) -> ta.Any:
1610
+ def _recv_obj(f: ta.IO, ty: ta.Any) -> ta.Any:
1522
1611
  d = f.read(4)
1523
1612
  if not d:
1524
1613
  return None
1525
1614
  if len(d) != 4:
1526
- raise Exception
1615
+ raise EOFError
1527
1616
 
1528
1617
  sz = struct.unpack('<I', d)[0]
1529
1618
  d = f.read(sz)
1530
- if not d:
1531
- raise Exception
1619
+ if len(d) != sz:
1620
+ raise EOFError
1532
1621
 
1533
1622
  j = json.loads(d.decode('utf-8'))
1534
1623
  return unmarshal_obj(j, ty)
@@ -1541,13 +1630,16 @@ def _remote_main() -> None:
1541
1630
  rt = pyremote_bootstrap_finalize() # noqa
1542
1631
 
1543
1632
  while True:
1544
- i = _recv_obj(rt.input, SubprocessCommand.Input)
1633
+ i = _recv_obj(rt.input, Command.Input)
1545
1634
  if i is None:
1546
1635
  break
1547
1636
 
1548
- o = SubprocessCommand()._execute(i) # noqa
1637
+ if isinstance(i, SubprocessCommand.Input):
1638
+ o = SubprocessCommand()._execute(i) # noqa
1639
+ else:
1640
+ raise TypeError(i)
1549
1641
 
1550
- _send_obj(rt.output, o)
1642
+ _send_obj(rt.output, o, Command.Output)
1551
1643
 
1552
1644
 
1553
1645
  ##
@@ -1634,8 +1726,13 @@ def _main() -> None:
1634
1726
  stdin = check_not_none(proc.stdin)
1635
1727
  stdout = check_not_none(proc.stdout)
1636
1728
 
1637
- res = PyremoteBootstrapDriver(remote_src).run(stdin, stdout)
1638
- print(res)
1729
+ res = PyremoteBootstrapDriver( # noqa
1730
+ remote_src,
1731
+ PyremoteBootstrapOptions(
1732
+ # debug=True,
1733
+ ),
1734
+ ).run(stdin, stdout)
1735
+ # print(res)
1639
1736
 
1640
1737
  #
1641
1738
 
@@ -1650,9 +1747,9 @@ def _main() -> None:
1650
1747
  capture_stdout=True,
1651
1748
  ),
1652
1749
  ]:
1653
- _send_obj(stdin, ci)
1750
+ _send_obj(stdin, ci, Command.Input)
1654
1751
 
1655
- o = _recv_obj(stdout, SubprocessCommand.Output)
1752
+ o = _recv_obj(stdout, Command.Output)
1656
1753
 
1657
1754
  print(o)
1658
1755
 
@@ -1,6 +1,10 @@
1
1
  #!/usr/bin/env python3
2
2
  # @omlish-amalg ./_manage.py
3
3
  # ruff: noqa: UP006 UP007
4
+ """
5
+ manage.py -s 'docker run -i python:3.12'
6
+ manage.py -qs 'ssh -i foo/bar foo@bar.baz' --python=python3.8
7
+ """
4
8
  import inspect
5
9
  import json
6
10
  import shlex
@@ -12,21 +16,59 @@ import typing as ta
12
16
  from omlish.lite.cached import cached_nullary
13
17
  from omlish.lite.check import check_not_none
14
18
  from omlish.lite.json import json_dumps_compact
19
+ from omlish.lite.marshal import PolymorphicObjMarshaler
20
+ from omlish.lite.marshal import get_obj_marshaler
15
21
  from omlish.lite.marshal import marshal_obj
22
+ from omlish.lite.marshal import register_opj_marshaler
16
23
  from omlish.lite.marshal import unmarshal_obj
17
24
  from omlish.lite.subprocesses import subprocess_maybe_shell_wrap_exec
18
25
 
19
26
  from ...pyremote import PyremoteBootstrapDriver
27
+ from ...pyremote import PyremoteBootstrapOptions
20
28
  from ...pyremote import pyremote_bootstrap_finalize
21
29
  from ...pyremote import pyremote_build_bootstrap_cmd
30
+ from .commands.base import Command
22
31
  from .commands.subprocess import SubprocessCommand
23
32
 
24
33
 
25
34
  ##
26
35
 
27
36
 
28
- def _send_obj(f: ta.IO, o: ta.Any) -> None:
29
- j = json_dumps_compact(marshal_obj(o))
37
+ _COMMAND_TYPES = {
38
+ 'subprocess': SubprocessCommand,
39
+ }
40
+
41
+
42
+ register_opj_marshaler(
43
+ Command.Input,
44
+ PolymorphicObjMarshaler.of([
45
+ PolymorphicObjMarshaler.Impl(
46
+ cty.Input,
47
+ k,
48
+ get_obj_marshaler(cty.Input),
49
+ )
50
+ for k, cty in _COMMAND_TYPES.items()
51
+ ]),
52
+ )
53
+
54
+ register_opj_marshaler(
55
+ Command.Output,
56
+ PolymorphicObjMarshaler.of([
57
+ PolymorphicObjMarshaler.Impl(
58
+ cty.Output,
59
+ k,
60
+ get_obj_marshaler(cty.Output),
61
+ )
62
+ for k, cty in _COMMAND_TYPES.items()
63
+ ]),
64
+ )
65
+
66
+
67
+ ##
68
+
69
+
70
+ def _send_obj(f: ta.IO, o: ta.Any, ty: ta.Any = None) -> None:
71
+ j = json_dumps_compact(marshal_obj(o, ty))
30
72
  d = j.encode('utf-8')
31
73
 
32
74
  f.write(struct.pack('<I', len(d)))
@@ -34,17 +76,17 @@ def _send_obj(f: ta.IO, o: ta.Any) -> None:
34
76
  f.flush()
35
77
 
36
78
 
37
- def _recv_obj(f: ta.IO, ty: type) -> ta.Any:
79
+ def _recv_obj(f: ta.IO, ty: ta.Any) -> ta.Any:
38
80
  d = f.read(4)
39
81
  if not d:
40
82
  return None
41
83
  if len(d) != 4:
42
- raise Exception
84
+ raise EOFError
43
85
 
44
86
  sz = struct.unpack('<I', d)[0]
45
87
  d = f.read(sz)
46
- if not d:
47
- raise Exception
88
+ if len(d) != sz:
89
+ raise EOFError
48
90
 
49
91
  j = json.loads(d.decode('utf-8'))
50
92
  return unmarshal_obj(j, ty)
@@ -57,13 +99,16 @@ def _remote_main() -> None:
57
99
  rt = pyremote_bootstrap_finalize() # noqa
58
100
 
59
101
  while True:
60
- i = _recv_obj(rt.input, SubprocessCommand.Input)
102
+ i = _recv_obj(rt.input, Command.Input)
61
103
  if i is None:
62
104
  break
63
105
 
64
- o = SubprocessCommand()._execute(i) # noqa
106
+ if isinstance(i, SubprocessCommand.Input):
107
+ o = SubprocessCommand()._execute(i) # noqa
108
+ else:
109
+ raise TypeError(i)
65
110
 
66
- _send_obj(rt.output, o)
111
+ _send_obj(rt.output, o, Command.Output)
67
112
 
68
113
 
69
114
  ##
@@ -150,8 +195,13 @@ def _main() -> None:
150
195
  stdin = check_not_none(proc.stdin)
151
196
  stdout = check_not_none(proc.stdout)
152
197
 
153
- res = PyremoteBootstrapDriver(remote_src).run(stdin, stdout)
154
- print(res)
198
+ res = PyremoteBootstrapDriver( # noqa
199
+ remote_src,
200
+ PyremoteBootstrapOptions(
201
+ # debug=True,
202
+ ),
203
+ ).run(stdin, stdout)
204
+ # print(res)
155
205
 
156
206
  #
157
207
 
@@ -166,9 +216,9 @@ def _main() -> None:
166
216
  capture_stdout=True,
167
217
  ),
168
218
  ]:
169
- _send_obj(stdin, ci)
219
+ _send_obj(stdin, ci, Command.Input)
170
220
 
171
- o = _recv_obj(stdout, SubprocessCommand.Output)
221
+ o = _recv_obj(stdout, Command.Output)
172
222
 
173
223
  print(o)
174
224