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

ominfra/manage/new/main.py

@@ -1,32 +1,74 @@
 #!/usr/bin/env python3
 # @omlish-amalg ./_manage.py
 # ruff: noqa: UP006 UP007
+"""
+manage.py -s 'docker run -i python:3.12'
+manage.py -qs 'ssh -i foo/bar foo@bar.baz' --python=python3.8
+"""
 import inspect
 import json
-import os
 import shlex
 import struct
 import subprocess
 import sys
 import typing as ta
 
+from omlish.lite.cached import cached_nullary
 from omlish.lite.check import check_not_none
 from omlish.lite.json import json_dumps_compact
+from omlish.lite.marshal import PolymorphicObjMarshaler
+from omlish.lite.marshal import get_obj_marshaler
 from omlish.lite.marshal import marshal_obj
+from omlish.lite.marshal import register_opj_marshaler
 from omlish.lite.marshal import unmarshal_obj
 from omlish.lite.subprocesses import subprocess_maybe_shell_wrap_exec
 
 from ...pyremote import PyremoteBootstrapDriver
+from ...pyremote import PyremoteBootstrapOptions
 from ...pyremote import pyremote_bootstrap_finalize
 from ...pyremote import pyremote_build_bootstrap_cmd
+from .commands.base import Command
 from .commands.subprocess import SubprocessCommand
 
 
 ##
 
 
-def _send_obj(f: ta.IO, o: ta.Any) -> None:
-    j = json_dumps_compact(marshal_obj(o))
+_COMMAND_TYPES = {
+    'subprocess': SubprocessCommand,
+}
+
+
+register_opj_marshaler(
+    Command.Input,
+    PolymorphicObjMarshaler.of([
+        PolymorphicObjMarshaler.Impl(
+            cty.Input,
+            k,
+            get_obj_marshaler(cty.Input),
+        )
+        for k, cty in _COMMAND_TYPES.items()
+    ]),
+)
+
+register_opj_marshaler(
+    Command.Output,
+    PolymorphicObjMarshaler.of([
+        PolymorphicObjMarshaler.Impl(
+            cty.Output,
+            k,
+            get_obj_marshaler(cty.Output),
+        )
+        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)))
@@ -34,33 +76,74 @@ def _send_obj(f: ta.IO, o: ta.Any) -> None:
     f.flush()
 
 
-def _recv_obj(f: ta.IO, ty: type) -> ta.Any:
+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 Exception
+        raise EOFError
 
     sz = struct.unpack('<I', d)[0]
     d = f.read(sz)
-    if not d:
-        raise Exception
+    if len(d) != sz:
+        raise EOFError
 
     j = json.loads(d.decode('utf-8'))
     return unmarshal_obj(j, ty)
 
 
+##
+
+
 def _remote_main() -> None:
     rt = pyremote_bootstrap_finalize()  # noqa
 
     while True:
-        i = _recv_obj(rt.input, SubprocessCommand.Input)
+        i = _recv_obj(rt.input, Command.Input)
         if i is None:
             break
 
-        o = SubprocessCommand()._execute(i)  # noqa
+        if isinstance(i, SubprocessCommand.Input):
+            o = SubprocessCommand()._execute(i)  # noqa
+        else:
+            raise TypeError(i)
+
+        _send_obj(rt.output, o, Command.Output)
+
+
+##
+
+
+@cached_nullary
+def _get_self_src() -> str:
+    return inspect.getsource(sys.modules[__name__])
+
+
+def _is_src_amalg(src: str) -> bool:
+    for l in src.splitlines():  # noqa
+        if l.startswith('# @omlish-amalg-output '):
+            return True
+    return False
+
+
+@cached_nullary
+def _is_self_amalg() -> bool:
+    return _is_src_amalg(_get_self_src())
 
-        _send_obj(rt.output, o)
+
+def _get_amalg_src(*, amalg_file: ta.Optional[str]) -> str:
+    if amalg_file is not None:
+        with open(amalg_file) as f:
+            return f.read()
+
+    if _is_self_amalg():
+        return _get_self_src()
+
+    import importlib.resources
+    return importlib.resources.read_text(__package__, '_manage.py')
+
+
+##
 
 
 def _main() -> None:
@@ -77,23 +160,7 @@ def _main() -> None:
 
     #
 
-    self_src = inspect.getsource(sys.modules[__name__])
-    self_src_lines = self_src.splitlines()
-    for l in self_src_lines:
-        if l.startswith('# @omlish-amalg-output '):
-            is_self_amalg = True
-            break
-    else:
-        is_self_amalg = False
-
-    if is_self_amalg:
-        amalg_src = self_src
-    else:
-        amalg_file = args._amalg_file  # noqa
-        if amalg_file is None:
-            amalg_file = os.path.join(os.path.dirname(__file__), '_manage.py')
-        with open(amalg_file) as f:
-            amalg_src = f.read()
+    amalg_src = _get_amalg_src(amalg_file=args._amalg_file)  # noqa
 
     #
 
@@ -128,8 +195,13 @@ def _main() -> None:
     stdin = check_not_none(proc.stdin)
     stdout = check_not_none(proc.stdout)
 
-    res = PyremoteBootstrapDriver(remote_src).run(stdin, stdout)
-    print(res)
+    res = PyremoteBootstrapDriver(  # noqa
+        remote_src,
+        PyremoteBootstrapOptions(
+            # debug=True,
+        ),
+    ).run(stdin, stdout)
+    # print(res)
 
     #
 
@@ -144,9 +216,9 @@ def _main() -> None:
             capture_stdout=True,
         ),
     ]:
-        _send_obj(stdin, ci)
+        _send_obj(stdin, ci, Command.Input)
 
-        o = _recv_obj(stdout, SubprocessCommand.Output)
+        o = _recv_obj(stdout, Command.Output)
 
         print(o)
 

ominfra/pyremote.py

@@ -5,7 +5,6 @@ Basically this: https://mitogen.networkgenomics.com/howitworks.html
 """
 import base64
 import dataclasses as dc
-import inspect
 import json
 import os
 import platform
@@ -13,7 +12,6 @@ import pwd
 import site
 import struct
 import sys
-import textwrap
 import typing as ta
 import zlib
 
@@ -21,11 +19,120 @@ import zlib
 ##
 
 
+@dc.dataclass(frozen=True)
+class PyremoteBootstrapOptions:
+    debug: bool = False
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class PyremoteEnvInfo:
+    sys_base_prefix: str
+    sys_byteorder: str
+    sys_defaultencoding: str
+    sys_exec_prefix: str
+    sys_executable: str
+    sys_implementation_name: str
+    sys_path: ta.List[str]
+    sys_platform: str
+    sys_prefix: str
+    sys_version: str
+    sys_version_info: ta.List[ta.Union[int, str]]
+
+    platform_architecture: ta.List[str]
+    platform_machine: str
+    platform_platform: str
+    platform_processor: str
+    platform_system: str
+    platform_release: str
+    platform_version: str
+
+    site_userbase: str
+
+    os_cwd: str
+    os_gid: int
+    os_loadavg: ta.List[float]
+    os_login: ta.Optional[str]
+    os_pgrp: int
+    os_pid: int
+    os_ppid: int
+    os_uid: int
+
+    pw_name: str
+    pw_uid: int
+    pw_gid: int
+    pw_gecos: str
+    pw_dir: str
+    pw_shell: str
+
+    env_path: ta.Optional[str]
+
+
+def _get_pyremote_env_info() -> PyremoteEnvInfo:
+    os_uid = os.getuid()
+
+    pw = pwd.getpwuid(os_uid)
+
+    os_login: ta.Optional[str]
+    try:
+        os_login = os.getlogin()
+    except OSError:
+        os_login = None
+
+    return PyremoteEnvInfo(
+        sys_base_prefix=sys.base_prefix,
+        sys_byteorder=sys.byteorder,
+        sys_defaultencoding=sys.getdefaultencoding(),
+        sys_exec_prefix=sys.exec_prefix,
+        sys_executable=sys.executable,
+        sys_implementation_name=sys.implementation.name,
+        sys_path=sys.path,
+        sys_platform=sys.platform,
+        sys_prefix=sys.prefix,
+        sys_version=sys.version,
+        sys_version_info=list(sys.version_info),
+
+        platform_architecture=list(platform.architecture()),
+        platform_machine=platform.machine(),
+        platform_platform=platform.platform(),
+        platform_processor=platform.processor(),
+        platform_system=platform.system(),
+        platform_release=platform.release(),
+        platform_version=platform.version(),
+
+        site_userbase=site.getuserbase(),
+
+        os_cwd=os.getcwd(),
+        os_gid=os.getgid(),
+        os_loadavg=list(os.getloadavg()),
+        os_login=os_login,
+        os_pgrp=os.getpgrp(),
+        os_pid=os.getpid(),
+        os_ppid=os.getppid(),
+        os_uid=os_uid,
+
+        pw_name=pw.pw_name,
+        pw_uid=pw.pw_uid,
+        pw_gid=pw.pw_gid,
+        pw_gecos=pw.pw_gecos,
+        pw_dir=pw.pw_dir,
+        pw_shell=pw.pw_shell,
+
+        env_path=os.environ.get('PATH'),
+    )
+
+
+##
+
+
 _PYREMOTE_BOOTSTRAP_INPUT_FD = 100
 _PYREMOTE_BOOTSTRAP_SRC_FD = 101
 
 _PYREMOTE_BOOTSTRAP_CHILD_PID_VAR = '_OPYR_CHILD_PID'
 _PYREMOTE_BOOTSTRAP_ARGV0_VAR = '_OPYR_ARGV0'
+_PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR = '_OPYR_OPTIONS_JSON'
 
 _PYREMOTE_BOOTSTRAP_ACK0 = b'OPYR000\n'
 _PYREMOTE_BOOTSTRAP_ACK1 = b'OPYR001\n'
@@ -87,7 +194,9 @@ def _pyremote_bootstrap_main(context_name: str) -> None:
 
         # Read main src from stdin
         main_z_len = struct.unpack('<I', os.read(0, 4))[0]
-        main_src = zlib.decompress(os.fdopen(0, 'rb').read(main_z_len))
+        if len(main_z := os.fdopen(0, 'rb').read(main_z_len)) != main_z_len:
+            raise EOFError
+        main_src = zlib.decompress(main_z)
 
         # Write both copies of main src. Must write to w0 (parent stdin) before w1 (copy pipe) as pipe will likely fill
         # and block and need to be drained by pyremote_bootstrap_finalize running in parent.
@@ -110,6 +219,8 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
     if any(c in context_name for c in '\'"'):
         raise NameError(context_name)
 
+    import inspect
+    import textwrap
     bs_src = textwrap.dedent(inspect.getsource(_pyremote_bootstrap_main))
 
     for gl in [
@@ -136,9 +247,6 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
     bs_z = zlib.compress(bs_src.encode('utf-8'))
     bs_z64 = base64.encodebytes(bs_z).replace(b'\n', b'')
 
-    def dq_repr(o: ta.Any) -> str:
-        return '"' + repr(o)[1:-1] + '"'
-
     stmts = [
         f'import {", ".join(_PYREMOTE_BOOTSTRAP_IMPORTS)}',
         f'exec(zlib.decompress(base64.decodebytes({bs_z64!r})))',
@@ -153,99 +261,83 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
 
 
 @dc.dataclass(frozen=True)
-class PyremoteEnvInfo:
-    sys_base_prefix: str
-    sys_byteorder: str
-    sys_defaultencoding: str
-    sys_exec_prefix: str
-    sys_executable: str
-    sys_implementation_name: str
-    sys_path: ta.List[str]
-    sys_platform: str
-    sys_prefix: str
-    sys_version: str
-    sys_version_info: ta.List[ta.Union[int, str]]
-
-    platform_architecture: ta.List[str]
-    platform_machine: str
-    platform_platform: str
-    platform_processor: str
-    platform_system: str
-    platform_release: str
-    platform_version: str
-
-    site_userbase: str
-
-    os_cwd: str
-    os_gid: int
-    os_loadavg: ta.List[float]
-    os_login: ta.Optional[str]
-    os_pgrp: int
-    os_pid: int
-    os_ppid: int
-    os_uid: int
-
-    pw_name: str
-    pw_uid: int
-    pw_gid: int
-    pw_gecos: str
-    pw_dir: str
-    pw_shell: str
-
-    env_path: ta.Optional[str]
+class PyremotePayloadRuntime:
+    input: ta.BinaryIO
+    output: ta.BinaryIO
+    main_src: str
+    options: PyremoteBootstrapOptions
+    env_info: PyremoteEnvInfo
 
 
-def _get_pyremote_env_info() -> PyremoteEnvInfo:
-    os_uid = os.getuid()
+def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
+    # If json options var is not present we need to do initial finalization
+    if _PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR not in os.environ:
+        # Read second copy of main src
+        r1 = os.fdopen(_PYREMOTE_BOOTSTRAP_SRC_FD, 'rb', 0)
+        main_src = r1.read().decode('utf-8')
+        r1.close()
+
+        # Reap boostrap child. Must be done after reading second copy of source because source may be too big to fit in
+        # a pipe at once.
+        os.waitpid(int(os.environ.pop(_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR)), 0)
+
+        # Read options
+        options_json_len = struct.unpack('<I', os.read(_PYREMOTE_BOOTSTRAP_INPUT_FD, 4))[0]
+        if len(options_json := os.read(_PYREMOTE_BOOTSTRAP_INPUT_FD, options_json_len)) != options_json_len:
+            raise EOFError
+        options = PyremoteBootstrapOptions(**json.loads(options_json.decode('utf-8')))
+
+        # If debugging, re-exec as file
+        if options.debug:
+            # Write temp source file
+            import tempfile
+            tfd, tfn = tempfile.mkstemp('-pyremote.py')
+            os.write(tfd, main_src.encode('utf-8'))
+            os.close(tfd)
+
+            # Set json options var
+            os.environ[_PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR] = options_json.decode('utf-8')
+
+            # Re-exec temp file
+            os.execl(os.environ[_PYREMOTE_BOOTSTRAP_ARGV0_VAR], sys.orig_argv[0], tfn)
 
-    pw = pwd.getpwuid(os_uid)
+    else:
+        # Load options json var
+        options_json_str = os.environ.pop(_PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR)
+        options = PyremoteBootstrapOptions(**json.loads(options_json_str))
 
-    os_login: ta.Optional[str]
-    try:
-        os_login = os.getlogin()
-    except OSError:
-        os_login = None
+        # Read temp source file
+        with open(sys.orig_argv[1]) as sf:
+            main_src = sf.read()
 
-    return PyremoteEnvInfo(
-        sys_base_prefix=sys.base_prefix,
-        sys_byteorder=sys.byteorder,
-        sys_defaultencoding=sys.getdefaultencoding(),
-        sys_exec_prefix=sys.exec_prefix,
-        sys_executable=sys.executable,
-        sys_implementation_name=sys.implementation.name,
-        sys_path=sys.path,
-        sys_platform=sys.platform,
-        sys_prefix=sys.prefix,
-        sys_version=sys.version,
-        sys_version_info=list(sys.version_info),
+    # Restore original argv0
+    sys.executable = os.environ.pop(_PYREMOTE_BOOTSTRAP_ARGV0_VAR)
 
-        platform_architecture=list(platform.architecture()),
-        platform_machine=platform.machine(),
-        platform_platform=platform.platform(),
-        platform_processor=platform.processor(),
-        platform_system=platform.system(),
-        platform_release=platform.release(),
-        platform_version=platform.version(),
+    # Write third ack
+    os.write(1, _PYREMOTE_BOOTSTRAP_ACK2)
 
-        site_userbase=site.getuserbase(),
+    # Write env info
+    env_info = _get_pyremote_env_info()
+    env_info_json = json.dumps(dc.asdict(env_info), indent=None, separators=(',', ':'))  # noqa
+    os.write(1, struct.pack('<I', len(env_info_json)))
+    os.write(1, env_info_json.encode('utf-8'))
 
-        os_cwd=os.getcwd(),
-        os_gid=os.getgid(),
-        os_loadavg=list(os.getloadavg()),
-        os_login=os_login,
-        os_pgrp=os.getpgrp(),
-        os_pid=os.getpid(),
-        os_ppid=os.getppid(),
-        os_uid=os_uid,
+    # Setup IO
+    input = os.fdopen(_PYREMOTE_BOOTSTRAP_INPUT_FD, 'rb', 0)  # noqa
+    output = os.fdopen(os.dup(1), 'wb', 0)  # noqa
+    os.dup2(nfd := os.open('/dev/null', os.O_WRONLY), 1)
+    os.close(nfd)
 
-        pw_name=pw.pw_name,
-        pw_uid=pw.pw_uid,
-        pw_gid=pw.pw_gid,
-        pw_gecos=pw.pw_gecos,
-        pw_dir=pw.pw_dir,
-        pw_shell=pw.pw_shell,
+    # Write fourth ack
+    output.write(_PYREMOTE_BOOTSTRAP_ACK3)
 
-        env_path=os.environ.get('PATH'),
+    # Return
+    return PyremotePayloadRuntime(
+        input=input,
+        output=output,
+        main_src=main_src,
+        options=options,
+        env_info=env_info,
     )
 
 
@@ -253,12 +345,15 @@ def _get_pyremote_env_info() -> PyremoteEnvInfo:
 
 
 class PyremoteBootstrapDriver:
-    def __init__(self, main_src: str) -> None:
+    def __init__(self, main_src: str, options: PyremoteBootstrapOptions = PyremoteBootstrapOptions()) -> None:
         super().__init__()
 
         self._main_src = main_src
         self._main_z = zlib.compress(main_src.encode('utf-8'))
 
+        self._options = options
+        self._options_json = json.dumps(dc.asdict(options), indent=None, separators=(',', ':')).encode('utf-8')  # noqa
+
     #
 
     @dc.dataclass(frozen=True)
@@ -278,7 +373,7 @@ class PyremoteBootstrapDriver:
         env_info: PyremoteEnvInfo
 
     def gen(self) -> ta.Generator[ta.Union[Read, Write], ta.Optional[bytes], Result]:
-        # Read first ack
+        # Read first ack (after fork)
         yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK0)
 
         # Read pid
@@ -289,8 +384,14 @@ class PyremoteBootstrapDriver:
         yield from self._write(struct.pack('<I', len(self._main_z)))
         yield from self._write(self._main_z)
 
-        # Read second and third ack
+        # Read second ack (after writing src copies)
         yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK1)
+
+        # Write options
+        yield from self._write(struct.pack('<I', len(self._options_json)))
+        yield from self._write(self._options_json)
+
+        # Read third ack (after reaping child process)
         yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK2)
 
         # Read env info
@@ -300,7 +401,7 @@ class PyremoteBootstrapDriver:
         env_info_json = d.decode('utf-8')
         env_info = PyremoteEnvInfo(**json.loads(env_info_json))
 
-        # Read fourth ack
+        # Read fourth ack (after finalization completed)
         yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK3)
 
         # Return
@@ -343,61 +444,11 @@ class PyremoteBootstrapDriver:
                 return e.value
 
             if isinstance(go, self.Read):
-                gi = stdout.read(go.sz)
+                if len(gi := stdout.read(go.sz)) != go.sz:
+                    raise EOFError
             elif isinstance(go, self.Write):
                 gi = None
                 stdin.write(go.d)
                 stdin.flush()
             else:
                 raise TypeError(go)
-
-
-##
-
-
-@dc.dataclass(frozen=True)
-class PyremotePayloadRuntime:
-    input: ta.BinaryIO
-    output: ta.BinaryIO
-    main_src: str
-    env_info: PyremoteEnvInfo
-
-
-def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
-    # Restore original argv0
-    sys.executable = os.environ.pop(_PYREMOTE_BOOTSTRAP_ARGV0_VAR)
-
-    # Read second copy of main src
-    r1 = os.fdopen(_PYREMOTE_BOOTSTRAP_SRC_FD, 'rb', 0)
-    main_src = r1.read().decode('utf-8')
-    r1.close()
-
-    # Reap boostrap child. Must be done after reading second copy of source because source may be too big to fit in a
-    # pipe at once.
-    os.waitpid(int(os.environ.pop(_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR)), 0)
-
-    # Write third ack
-    os.write(1, _PYREMOTE_BOOTSTRAP_ACK2)
-
-    # Write env info
-    env_info = _get_pyremote_env_info()
-    env_info_json = json.dumps(dc.asdict(env_info), indent=None, separators=(',', ':'))  # noqa
-    os.write(1, struct.pack('<I', len(env_info_json)))
-    os.write(1, env_info_json.encode('utf-8'))
-
-    # Setup IO
-    input = os.fdopen(_PYREMOTE_BOOTSTRAP_INPUT_FD, 'rb', 0)  # noqa
-    output = os.fdopen(os.dup(1), 'wb', 0)  # noqa
-    os.dup2(nfd := os.open('/dev/null', os.O_WRONLY), 1)
-    os.close(nfd)
-
-    # Write fourth ack
-    output.write(_PYREMOTE_BOOTSTRAP_ACK3)
-
-    # Return
-    return PyremotePayloadRuntime(
-        input=input,
-        output=output,
-        main_src=main_src,
-        env_info=env_info,
-    )