ominfra 0.0.0.dev133__py3-none-any.whl → 0.0.0.dev134__py3-none-any.whl

ominfra/pyremote.py

@@ -17,9 +17,6 @@ import textwrap
 import typing as ta
 import zlib
 
-from omlish.lite.check import check_isinstance
-from omlish.lite.check import check_none
-
 
 ##
 
@@ -32,8 +29,8 @@ _PYREMOTE_BOOTSTRAP_ARGV0_VAR = '_OPYR_ARGV0'
 
 _PYREMOTE_BOOTSTRAP_ACK0 = b'OPYR000\n'
 _PYREMOTE_BOOTSTRAP_ACK1 = b'OPYR001\n'
-_PYREMOTE_BOOTSTRAP_ACK2 = b'OPYR001\n'
-_PYREMOTE_BOOTSTRAP_ACK3 = b'OPYR001\n'
+_PYREMOTE_BOOTSTRAP_ACK2 = b'OPYR002\n'
+_PYREMOTE_BOOTSTRAP_ACK3 = b'OPYR003\n'
 
 _PYREMOTE_BOOTSTRAP_PROC_TITLE_FMT = '(pyremote:%s)'
 
@@ -256,6 +253,8 @@ class PyremoteBootstrapDriver:
         self._main_src = main_src
         self._main_z = zlib.compress(main_src.encode('utf-8'))
 
+    #
+
     @dc.dataclass(frozen=True)
     class Read:
         sz: int
@@ -272,7 +271,7 @@ class PyremoteBootstrapDriver:
         pid: int
         env_info: PyremoteEnvInfo
 
-    def __call__(self) -> ta.Generator[ta.Union[Read, Write], ta.Optional[bytes], Result]:
+    def gen(self) -> ta.Generator[ta.Union[Read, Write], ta.Optional[bytes], Result]:
         # Read first ack
         yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK0)
 
@@ -281,8 +280,8 @@ class PyremoteBootstrapDriver:
         pid = struct.unpack('<Q', d)[0]
 
         # Write main src
-        check_none((yield self.Write(struct.pack('<I', len(self._main_z)))))
-        check_none((yield self.Write(self._main_z)))
+        yield from self._write(struct.pack('<I', len(self._main_z)))
+        yield from self._write(self._main_z)
 
         # Read second and third ack
         yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK1)
@@ -305,7 +304,9 @@ class PyremoteBootstrapDriver:
         )
 
     def _read(self, sz: int) -> ta.Generator[Read, bytes, bytes]:
-        d = check_isinstance((yield self.Read(sz)), bytes)
+        d = yield self.Read(sz)
+        if not isinstance(d, bytes):
+            raise self.ProtocolError(f'Expected bytes after read, got {d!r}')
         if len(d) != sz:
             raise self.ProtocolError(f'Read {len(d)} bytes, expected {sz}')
         return d
@@ -315,17 +316,47 @@ class PyremoteBootstrapDriver:
         if d != e:
             raise self.ProtocolError(f'Read {d!r}, expected {e!r}')
 
+    def _write(self, d: bytes) -> ta.Generator[Write, ta.Optional[bytes], None]:
+        i = yield self.Write(d)
+        if i is not None:
+            raise self.ProtocolError('Unexpected input after write')
+
+    #
+
+    def run(self, stdin: ta.IO, stdout: ta.IO) -> Result:
+        gen = self.gen()
+
+        gi: bytes | None = None
+        while True:
+            try:
+                if gi is not None:
+                    go = gen.send(gi)
+                else:
+                    go = next(gen)
+            except StopIteration as e:
+                return e.value
+
+            if isinstance(go, self.Read):
+                gi = stdout.read(go.sz)
+            elif isinstance(go, self.Write):
+                gi = None
+                stdin.write(go.d)
+                stdin.flush()
+            else:
+                raise TypeError(go)
+
 
 ##
 
 
 @dc.dataclass(frozen=True)
-class PyremoteBootstrapPayloadRuntime:
+class PyremotePayloadRuntime:
     input: ta.BinaryIO
     main_src: str
+    env_info: PyremoteEnvInfo
 
 
-def pyremote_bootstrap_finalize() -> PyremoteBootstrapPayloadRuntime:
+def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
     # Restore original argv0
     sys.executable = os.environ.pop(_PYREMOTE_BOOTSTRAP_ARGV0_VAR)
 
@@ -351,7 +382,8 @@ def pyremote_bootstrap_finalize() -> PyremoteBootstrapPayloadRuntime:
     os.write(1, _PYREMOTE_BOOTSTRAP_ACK3)
 
     # Return
-    return PyremoteBootstrapPayloadRuntime(
+    return PyremotePayloadRuntime(
         input=os.fdopen(_PYREMOTE_BOOTSTRAP_COMM_FD, 'rb', 0),
         main_src=main_src,
+        env_info=env_info,
     )

{ominfra-0.0.0.dev133.dist-info → ominfra-0.0.0.dev134.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ominfra
-Version: 0.0.0.dev133
+Version: 0.0.0.dev134
 Summary: ominfra
 Author: wrmsr
 License: BSD-3-Clause
@@ -12,8 +12,8 @@ Classifier: Operating System :: OS Independent
 Classifier: Operating System :: POSIX
 Requires-Python: >=3.12
 License-File: LICENSE
-Requires-Dist: omdev==0.0.0.dev133
-Requires-Dist: omlish==0.0.0.dev133
+Requires-Dist: omdev==0.0.0.dev134
+Requires-Dist: omlish==0.0.0.dev134
 Provides-Extra: all
 Requires-Dist: paramiko~=3.5; extra == "all"
 Requires-Dist: asyncssh~=2.18; extra == "all"

{ominfra-0.0.0.dev133.dist-info → ominfra-0.0.0.dev134.dist-info}/RECORD

@@ -3,7 +3,7 @@ ominfra/__about__.py,sha256=6i1AoruFYQCd-PyhhbDQDWY2d1tiQu9nkwWr-fXAqfY,705
 ominfra/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ominfra/cmds.py,sha256=E0AfnvEmnKntXWvmLW5L05_NeDpBET1VBXn7vV6EwBQ,2083
 ominfra/configs.py,sha256=8aU1Qmbr-qjaE2iP3gAbA2SWJYMPZ-uGK007L01PoOI,1727
-ominfra/pyremote.py,sha256=U7OYc4iF1HGRcESdq_KhoS1RqTWHk4bq8yFRORFX4O0,9264
+ominfra/pyremote.py,sha256=nqzmMKZvUMiw2y65v8VVjx832sKfSB7jBOkP8SQBmGI,10121
 ominfra/ssh.py,sha256=jQpc4WvkMckIfk4vILda8zFaeharRqc_6wxW50b0OjQ,5431
 ominfra/threadworkers.py,sha256=oX4ubZn7h932saXpRIJu2MNhBExgGGMuGhdXarZxLJw,4948
 ominfra/clouds/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -100,9 +100,9 @@ ominfra/tailscale/api.py,sha256=C5-t_b6jZXUWcy5k8bXm7CFnk73pSdrlMOgGDeGVrpw,1370
 ominfra/tailscale/cli.py,sha256=DSGp4hn5xwOW-l_u_InKlSF6kIobxtUtVssf_73STs0,3567
 ominfra/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ominfra/tools/listresources.py,sha256=4qVg5txsb10EHhvqXXeM6gJ2jx9LbroEnPydDv1uXs0,6176
-ominfra-0.0.0.dev133.dist-info/LICENSE,sha256=B_hVtavaA8zCYDW99DYdcpDLKz1n3BBRjZrcbv8uG8c,1451
-ominfra-0.0.0.dev133.dist-info/METADATA,sha256=_KW6-EvKpnZ8VyQPHOE1ueDkQ8zjfd2CJ-VAl_WsIbY,731
-ominfra-0.0.0.dev133.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-ominfra-0.0.0.dev133.dist-info/entry_points.txt,sha256=kgecQ2MgGrM9qK744BoKS3tMesaC3yjLnl9pa5CRczg,37
-ominfra-0.0.0.dev133.dist-info/top_level.txt,sha256=E-b2OHkk_AOBLXHYZQ2EOFKl-_6uOGd8EjeG-Zy6h_w,8
-ominfra-0.0.0.dev133.dist-info/RECORD,,
+ominfra-0.0.0.dev134.dist-info/LICENSE,sha256=B_hVtavaA8zCYDW99DYdcpDLKz1n3BBRjZrcbv8uG8c,1451
+ominfra-0.0.0.dev134.dist-info/METADATA,sha256=pJ8EYFED29tv1T8wzhlFFm2rKeNXXS2c4Iw8hjFEvSA,731
+ominfra-0.0.0.dev134.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+ominfra-0.0.0.dev134.dist-info/entry_points.txt,sha256=kgecQ2MgGrM9qK744BoKS3tMesaC3yjLnl9pa5CRczg,37
+ominfra-0.0.0.dev134.dist-info/top_level.txt,sha256=E-b2OHkk_AOBLXHYZQ2EOFKl-_6uOGd8EjeG-Zy6h_w,8
+ominfra-0.0.0.dev134.dist-info/RECORD,,