ominfra 0.0.0.dev141__py3-none-any.whl → 0.0.0.dev142__py3-none-any.whl

ominfra/scripts/manage.py

@@ -59,6 +59,9 @@ CallableT = ta.TypeVar('CallableT', bound=ta.Callable)
 # ../../omlish/lite/check.py
 SizedT = ta.TypeVar('SizedT', bound=ta.Sized)
 
+# ../../omlish/lite/subprocesses.py
+SubprocessChannelOption = ta.Literal['pipe', 'stdout', 'devnull']
+
 
 ########################################
 # ../commands/base.py
@@ -251,12 +254,14 @@ def _pyremote_bootstrap_main(context_name: str) -> None:
             os.close(f)
 
         # Save vars
-        os.environ[_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR] = str(cp)
-        os.environ[_PYREMOTE_BOOTSTRAP_ARGV0_VAR] = sys.executable
-        os.environ[_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR] = context_name
+        env = os.environ
+        exe = sys.executable
+        env[_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR] = str(cp)
+        env[_PYREMOTE_BOOTSTRAP_ARGV0_VAR] = exe
+        env[_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR] = context_name
 
         # Start repl reading stdin from r0
-        os.execl(sys.executable, sys.executable + (_PYREMOTE_BOOTSTRAP_PROC_TITLE_FMT % (context_name,)))
+        os.execl(exe, exe + (_PYREMOTE_BOOTSTRAP_PROC_TITLE_FMT % (context_name,)))
 
     else:
         # Child process
@@ -320,12 +325,12 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
         if cl.strip()
     )
 
-    bs_z = zlib.compress(bs_src.encode('utf-8'))
-    bs_z64 = base64.encodebytes(bs_z).replace(b'\n', b'')
+    bs_z = zlib.compress(bs_src.encode('utf-8'), 9)
+    bs_z85 = base64.b85encode(bs_z).replace(b'\n', b'')
 
     stmts = [
         f'import {", ".join(_PYREMOTE_BOOTSTRAP_IMPORTS)}',
-        f'exec(zlib.decompress(base64.decodebytes({bs_z64!r})))',
+        f'exec(zlib.decompress(base64.b85decode({bs_z85!r})))',
         f'_pyremote_bootstrap_main({context_name!r})',
     ]
 
@@ -514,7 +519,7 @@ class PyremoteBootstrapDriver:
 
     #
 
-    def run(self, stdin: ta.IO, stdout: ta.IO) -> Result:
+    def run(self, input: ta.IO, output: ta.IO) -> Result:  # noqa
         gen = self.gen()
 
         gi: ta.Optional[bytes] = None
@@ -528,12 +533,12 @@ class PyremoteBootstrapDriver:
                 return e.value
 
             if isinstance(go, self.Read):
-                if len(gi := stdout.read(go.sz)) != go.sz:
+                if len(gi := input.read(go.sz)) != go.sz:
                     raise EOFError
             elif isinstance(go, self.Write):
                 gi = None
-                stdin.write(go.d)
-                stdin.flush()
+                output.write(go.d)
+                output.flush()
             else:
                 raise TypeError(go)
 
@@ -1286,120 +1291,144 @@ _OBJ_MARSHALER_GENERIC_ITERABLE_TYPES: ta.Dict[ta.Any, type] = {
 }
 
 
-def _make_obj_marshaler(
-        ty: ta.Any,
-        rec: ta.Callable[[ta.Any], ObjMarshaler],
-        *,
-        nonstrict_dataclasses: bool = False,
-) -> ObjMarshaler:
-    if isinstance(ty, type):
-        if abc.ABC in ty.__bases__:
-            return PolymorphicObjMarshaler.of([  # type: ignore
-                PolymorphicObjMarshaler.Impl(
-                    ity,
-                    ity.__qualname__,
-                    rec(ity),
-                )
-                for ity in deep_subclasses(ty)
-                if abc.ABC not in ity.__bases__
-            ])
-
-        if issubclass(ty, enum.Enum):
-            return EnumObjMarshaler(ty)
-
-        if dc.is_dataclass(ty):
-            return DataclassObjMarshaler(
-                ty,
-                {f.name: rec(f.type) for f in dc.fields(ty)},
-                nonstrict=nonstrict_dataclasses,
-            )
+##
 
-    if is_generic_alias(ty):
-        try:
-            mt = _OBJ_MARSHALER_GENERIC_MAPPING_TYPES[ta.get_origin(ty)]
-        except KeyError:
-            pass
-        else:
-            k, v = ta.get_args(ty)
-            return MappingObjMarshaler(mt, rec(k), rec(v))
 
-        try:
-            st = _OBJ_MARSHALER_GENERIC_ITERABLE_TYPES[ta.get_origin(ty)]
-        except KeyError:
-            pass
-        else:
-            [e] = ta.get_args(ty)
-            return IterableObjMarshaler(st, rec(e))
+class ObjMarshalerManager:
+    def __init__(
+            self,
+            *,
+            default_obj_marshalers: ta.Dict[ta.Any, ObjMarshaler] = _DEFAULT_OBJ_MARSHALERS,  # noqa
+            generic_mapping_types: ta.Dict[ta.Any, type] = _OBJ_MARSHALER_GENERIC_MAPPING_TYPES,  # noqa
+            generic_iterable_types: ta.Dict[ta.Any, type] = _OBJ_MARSHALER_GENERIC_ITERABLE_TYPES,  # noqa
+    ) -> None:
+        super().__init__()
 
-        if is_union_alias(ty):
-            return OptionalObjMarshaler(rec(get_optional_alias_arg(ty)))
+        self._obj_marshalers = dict(default_obj_marshalers)
+        self._generic_mapping_types = generic_mapping_types
+        self._generic_iterable_types = generic_iterable_types
 
-    raise TypeError(ty)
+        self._lock = threading.RLock()
+        self._marshalers: ta.Dict[ta.Any, ObjMarshaler] = dict(_DEFAULT_OBJ_MARSHALERS)
+        self._proxies: ta.Dict[ta.Any, ProxyObjMarshaler] = {}
 
+    #
 
-##
+    def make_obj_marshaler(
+            self,
+            ty: ta.Any,
+            rec: ta.Callable[[ta.Any], ObjMarshaler],
+            *,
+            nonstrict_dataclasses: bool = False,
+    ) -> ObjMarshaler:
+        if isinstance(ty, type):
+            if abc.ABC in ty.__bases__:
+                return PolymorphicObjMarshaler.of([  # type: ignore
+                    PolymorphicObjMarshaler.Impl(
+                        ity,
+                        ity.__qualname__,
+                        rec(ity),
+                    )
+                    for ity in deep_subclasses(ty)
+                    if abc.ABC not in ity.__bases__
+                ])
+
+            if issubclass(ty, enum.Enum):
+                return EnumObjMarshaler(ty)
+
+            if dc.is_dataclass(ty):
+                return DataclassObjMarshaler(
+                    ty,
+                    {f.name: rec(f.type) for f in dc.fields(ty)},
+                    nonstrict=nonstrict_dataclasses,
+                )
 
+        if is_generic_alias(ty):
+            try:
+                mt = self._generic_mapping_types[ta.get_origin(ty)]
+            except KeyError:
+                pass
+            else:
+                k, v = ta.get_args(ty)
+                return MappingObjMarshaler(mt, rec(k), rec(v))
 
-_OBJ_MARSHALERS_LOCK = threading.RLock()
+            try:
+                st = self._generic_iterable_types[ta.get_origin(ty)]
+            except KeyError:
+                pass
+            else:
+                [e] = ta.get_args(ty)
+                return IterableObjMarshaler(st, rec(e))
 
-_OBJ_MARSHALERS: ta.Dict[ta.Any, ObjMarshaler] = dict(_DEFAULT_OBJ_MARSHALERS)
+            if is_union_alias(ty):
+                return OptionalObjMarshaler(rec(get_optional_alias_arg(ty)))
 
-_OBJ_MARSHALER_PROXIES: ta.Dict[ta.Any, ProxyObjMarshaler] = {}
+        raise TypeError(ty)
 
+    #
 
-def register_opj_marshaler(ty: ta.Any, m: ObjMarshaler) -> None:
-    with _OBJ_MARSHALERS_LOCK:
-        if ty in _OBJ_MARSHALERS:
-            raise KeyError(ty)
-        _OBJ_MARSHALERS[ty] = m
+    def register_opj_marshaler(self, ty: ta.Any, m: ObjMarshaler) -> None:
+        with self._lock:
+            if ty in self._obj_marshalers:
+                raise KeyError(ty)
+            self._obj_marshalers[ty] = m
 
+    def get_obj_marshaler(
+            self,
+            ty: ta.Any,
+            *,
+            no_cache: bool = False,
+            **kwargs: ta.Any,
+    ) -> ObjMarshaler:
+        with self._lock:
+            if not no_cache:
+                try:
+                    return self._obj_marshalers[ty]
+                except KeyError:
+                    pass
 
-def get_obj_marshaler(
-        ty: ta.Any,
-        *,
-        no_cache: bool = False,
-        **kwargs: ta.Any,
-) -> ObjMarshaler:
-    with _OBJ_MARSHALERS_LOCK:
-        if not no_cache:
             try:
-                return _OBJ_MARSHALERS[ty]
+                return self._proxies[ty]
             except KeyError:
                 pass
 
-        try:
-            return _OBJ_MARSHALER_PROXIES[ty]
-        except KeyError:
-            pass
-
-        rec = functools.partial(
-            get_obj_marshaler,
-            no_cache=no_cache,
-            **kwargs,
-        )
+            rec = functools.partial(
+                self.get_obj_marshaler,
+                no_cache=no_cache,
+                **kwargs,
+            )
 
-        p = ProxyObjMarshaler()
-        _OBJ_MARSHALER_PROXIES[ty] = p
-        try:
-            m = _make_obj_marshaler(ty, rec, **kwargs)
-        finally:
-            del _OBJ_MARSHALER_PROXIES[ty]
-        p.m = m
+            p = ProxyObjMarshaler()
+            self._proxies[ty] = p
+            try:
+                m = self.make_obj_marshaler(ty, rec, **kwargs)
+            finally:
+                del self._proxies[ty]
+            p.m = m
 
-        if not no_cache:
-            _OBJ_MARSHALERS[ty] = m
-        return m
+            if not no_cache:
+                self._obj_marshalers[ty] = m
+            return m
+
+    #
+
+    def marshal_obj(self, o: ta.Any, ty: ta.Any = None) -> ta.Any:
+        return self.get_obj_marshaler(ty if ty is not None else type(o)).marshal(o)
+
+    def unmarshal_obj(self, o: ta.Any, ty: ta.Union[ta.Type[T], ta.Any]) -> T:
+        return self.get_obj_marshaler(ty).unmarshal(o)
 
 
 ##
 
 
-def marshal_obj(o: ta.Any, ty: ta.Any = None) -> ta.Any:
-    return get_obj_marshaler(ty if ty is not None else type(o)).marshal(o)
+OBJ_MARSHALER_MANAGER = ObjMarshalerManager()
 
+register_opj_marshaler = OBJ_MARSHALER_MANAGER.register_opj_marshaler
+get_obj_marshaler = OBJ_MARSHALER_MANAGER.get_obj_marshaler
 
-def unmarshal_obj(o: ta.Any, ty: ta.Union[ta.Type[T], ta.Any]) -> T:
-    return get_obj_marshaler(ty).unmarshal(o)
+marshal_obj = OBJ_MARSHALER_MANAGER.marshal_obj
+unmarshal_obj = OBJ_MARSHALER_MANAGER.unmarshal_obj
 
 
 ########################################
@@ -1419,6 +1448,48 @@ def check_runtime_version() -> None:
         raise OSError(f'Requires python {REQUIRED_PYTHON_VERSION}, got {sys.version_info} from {sys.executable}')  # noqa
 
 
+########################################
+# ../protocol.py
+
+
+class Channel:
+    def __init__(
+            self,
+            input: ta.IO,  # noqa
+            output: ta.IO,
+            *,
+            msh: ObjMarshalerManager = OBJ_MARSHALER_MANAGER,
+    ) -> None:
+        super().__init__()
+
+        self._input = input
+        self._output = output
+        self._msh = msh
+
+    def send_obj(self, o: ta.Any, ty: ta.Any = None) -> None:
+        j = json_dumps_compact(self._msh.marshal_obj(o, ty))
+        d = j.encode('utf-8')
+
+        self._output.write(struct.pack('<I', len(d)))
+        self._output.write(d)
+        self._output.flush()
+
+    def recv_obj(self, ty: ta.Any) -> ta.Any:
+        d = self._input.read(4)
+        if not d:
+            return None
+        if len(d) != 4:
+            raise EOFError
+
+        sz = struct.unpack('<I', d)[0]
+        d = self._input.read(sz)
+        if len(d) != sz:
+            raise EOFError
+
+        j = json.loads(d.decode('utf-8'))
+        return self._msh.unmarshal_obj(j, ty)
+
+
 ########################################
 # ../../../omlish/lite/subprocesses.py
 
@@ -1426,6 +1497,16 @@ def check_runtime_version() -> None:
 ##
 
 
+SUBPROCESS_CHANNEL_OPTION_VALUES: ta.Mapping[SubprocessChannelOption, int] = {
+    'pipe': subprocess.PIPE,
+    'stdout': subprocess.STDOUT,
+    'devnull': subprocess.DEVNULL,
+}
+
+
+##
+
+
 _SUBPROCESS_SHELL_WRAP_EXECS = False
 
 
@@ -1552,21 +1633,21 @@ def subprocess_close(
 
 @dc.dataclass(frozen=True)
 class SubprocessCommand(Command['SubprocessCommand.Output']):
-    args: ta.Sequence[str]
+    cmd: ta.Sequence[str]
 
     shell: bool = False
     cwd: ta.Optional[str] = None
     env: ta.Optional[ta.Mapping[str, str]] = None
 
-    capture_stdout: bool = False
-    capture_stderr: bool = False
+    stdout: str = 'pipe'  # SubprocessChannelOption
+    stderr: str = 'pipe'  # SubprocessChannelOption
 
     input: ta.Optional[bytes] = None
     timeout: ta.Optional[float] = None
 
     def __post_init__(self) -> None:
-        if isinstance(self.args, str):
-            raise TypeError(self.args)
+        if isinstance(self.cmd, str):
+            raise TypeError(self.cmd)
 
     @dc.dataclass(frozen=True)
     class Output(Command.Output):
@@ -1584,24 +1665,23 @@ class SubprocessCommand(Command['SubprocessCommand.Output']):
 
 class SubprocessCommandExecutor(CommandExecutor[SubprocessCommand, SubprocessCommand.Output]):
     def execute(self, inp: SubprocessCommand) -> SubprocessCommand.Output:
-        proc = subprocess.Popen(
-            subprocess_maybe_shell_wrap_exec(*inp.args),
+        with subprocess.Popen(
+            subprocess_maybe_shell_wrap_exec(*inp.cmd),
 
             shell=inp.shell,
             cwd=inp.cwd,
             env={**os.environ, **(inp.env or {})},
 
             stdin=subprocess.PIPE if inp.input is not None else None,
-            stdout=subprocess.PIPE if inp.capture_stdout else None,
-            stderr=subprocess.PIPE if inp.capture_stderr else None,
-        )
-
-        start_time = time.time()
-        stdout, stderr = proc.communicate(
-            input=inp.input,
-            timeout=inp.timeout,
-        )
-        end_time = time.time()
+            stdout=SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, inp.stdout)],
+            stderr=SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, inp.stderr)],
+        ) as proc:
+            start_time = time.time()
+            stdout, stderr = proc.communicate(
+                input=inp.input,
+                timeout=inp.timeout,
+            )
+            end_time = time.time()
 
         return SubprocessCommand.Output(
             rc=proc.returncode,
@@ -1628,7 +1708,7 @@ class PySpawner:
             shell: ta.Optional[str] = None,
             shell_quote: bool = False,
             python: str = DEFAULT_PYTHON,
-            stderr: ta.Optional[ta.Literal['pipe', 'stdout', 'devnull']] = None,
+            stderr: ta.Optional[SubprocessChannelOption] = None,
     ) -> None:
         super().__init__()
 
@@ -1663,12 +1743,6 @@ class PySpawner:
 
     #
 
-    _STDERR_KWARG_MAP: ta.Mapping[str, int] = {
-        'pipe': subprocess.PIPE,
-        'stdout': subprocess.STDOUT,
-        'devnull': subprocess.DEVNULL,
-    }
-
     @dc.dataclass(frozen=True)
     class Spawned:
         stdin: ta.IO
@@ -1688,7 +1762,7 @@ class PySpawner:
             shell=pc.shell,
             stdin=subprocess.PIPE,
             stdout=subprocess.PIPE,
-            stderr=self._STDERR_KWARG_MAP[self._stderr] if self._stderr is not None else None,
+            stderr=SUBPROCESS_CHANNEL_OPTION_VALUES[self._stderr] if self._stderr is not None else None,
         ) as proc:
             stdin = check_not_none(proc.stdin)
             stdout = check_not_none(proc.stdout)
@@ -1721,51 +1795,28 @@ _COMMAND_TYPES = {
 }
 
 
-@static_init
-def _register_command_marshaling() -> None:
+##
+
+
+def register_command_marshaling(msh: ObjMarshalerManager) -> None:
     for fn in [
         lambda c: c,
         lambda c: c.Output,
     ]:
-        register_opj_marshaler(
+        msh.register_opj_marshaler(
             fn(Command),
             PolymorphicObjMarshaler.of([
                 PolymorphicObjMarshaler.Impl(
                     fn(cty),
                     k,
-                    get_obj_marshaler(fn(cty)),
+                    msh.get_obj_marshaler(fn(cty)),
                 )
                 for k, cty in _COMMAND_TYPES.items()
             ]),
         )
 
 
-##
-
-
-def _send_obj(f: ta.IO, o: ta.Any, ty: ta.Any = None) -> None:
-    j = json_dumps_compact(marshal_obj(o, ty))
-    d = j.encode('utf-8')
-
-    f.write(struct.pack('<I', len(d)))
-    f.write(d)
-    f.flush()
-
-
-def _recv_obj(f: ta.IO, ty: ta.Any) -> ta.Any:
-    d = f.read(4)
-    if not d:
-        return None
-    if len(d) != 4:
-        raise EOFError
-
-    sz = struct.unpack('<I', d)[0]
-    d = f.read(sz)
-    if len(d) != sz:
-        raise EOFError
-
-    j = json.loads(d.decode('utf-8'))
-    return unmarshal_obj(j, ty)
+register_command_marshaling(OBJ_MARSHALER_MANAGER)
 
 
 ##
@@ -1773,9 +1824,10 @@ def _recv_obj(f: ta.IO, ty: ta.Any) -> ta.Any:
 
 def _remote_main() -> None:
     rt = pyremote_bootstrap_finalize()  # noqa
+    chan = Channel(rt.input, rt.output)
 
     while True:
-        i = _recv_obj(rt.input, Command)
+        i = chan.recv_obj(Command)
         if i is None:
             break
 
@@ -1784,7 +1836,7 @@ def _remote_main() -> None:
         else:
             raise TypeError(i)
 
-        _send_obj(rt.output, o, Command.Output)
+        chan.send_obj(o, Command.Output)
 
 
 ##
@@ -1795,11 +1847,13 @@ def _main() -> None:
 
     parser = argparse.ArgumentParser()
 
+    parser.add_argument('--_payload-file')
+
     parser.add_argument('-s', '--shell')
     parser.add_argument('-q', '--shell-quote', action='store_true')
     parser.add_argument('--python', default='python3')
+
     parser.add_argument('--debug', action='store_true')
-    parser.add_argument('--_payload-file')
 
     args = parser.parse_args()
 
@@ -1807,8 +1861,6 @@ def _main() -> None:
 
     payload_src = get_payload_src(file=args._payload_file)  # noqa
 
-    #
-
     remote_src = '\n\n'.join([
         '__name__ = "__remote__"',
         payload_src,
@@ -1817,41 +1869,32 @@ def _main() -> None:
 
     #
 
-    bs_src = pyremote_build_bootstrap_cmd(__package__ or 'manage')
-
-    #
-
     spawner = PySpawner(
-        bs_src,
+        pyremote_build_bootstrap_cmd(__package__ or 'manage'),
         shell=args.shell,
         shell_quote=args.shell_quote,
         python=args.python,
     )
+
     with spawner.spawn() as proc:
         res = PyremoteBootstrapDriver(  # noqa
             remote_src,
             PyremoteBootstrapOptions(
                 debug=args.debug,
             ),
-        ).run(proc.stdin, proc.stdout)
-        # print(res)
+        ).run(proc.stdout, proc.stdin)
+
+        chan = Channel(proc.stdout, proc.stdin)
 
         #
 
         for ci in [
-            SubprocessCommand(
-                args=['python3', '-'],
-                input=b'print(1)\n',
-                capture_stdout=True,
-            ),
-            SubprocessCommand(
-                args=['uname'],
-                capture_stdout=True,
-            ),
+            SubprocessCommand(['python3', '-'], input=b'print(1)\n'),
+            SubprocessCommand(['uname']),
         ]:
-            _send_obj(proc.stdin, ci, Command)
+            chan.send_obj(ci, Command)
 
-            o = _recv_obj(proc.stdout, Command.Output)
+            o = chan.recv_obj(Command.Output)
 
             print(o)