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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -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