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

ominfra/{manage/new/_manage.py → scripts/manage.py}

@@ -2,11 +2,11 @@
 # noinspection DuplicatedCode
 # @omlish-lite
 # @omlish-script
-# @omlish-amalg-output main.py
+# @omlish-amalg-output ../manage/main.py
 # ruff: noqa: N802 UP006 UP007 UP036
 """
 manage.py -s 'docker run -i python:3.12'
-manage.py -qs 'ssh -i foo/bar foo@bar.baz' --python=python3.8
+manage.py -s 'ssh -i /foo/bar.pem foo@bar.baz' -q --python=python3.8
 """
 import abc
 import base64
@@ -49,13 +49,14 @@ if sys.version_info < (3, 8):
 
 
 # commands/base.py
-CommandInputT = ta.TypeVar('CommandInputT', bound='Command.Input')
+CommandT = ta.TypeVar('CommandT', bound='Command')
 CommandOutputT = ta.TypeVar('CommandOutputT', bound='Command.Output')
 
-# ../../../omlish/lite/cached.py
+# ../../omlish/lite/cached.py
 T = ta.TypeVar('T')
+CallableT = ta.TypeVar('CallableT', bound=ta.Callable)
 
-# ../../../omlish/lite/check.py
+# ../../omlish/lite/check.py
 SizedT = ta.TypeVar('SizedT', bound=ta.Sized)
 
 
@@ -66,22 +67,24 @@ SizedT = ta.TypeVar('SizedT', bound=ta.Sized)
 ##
 
 
-class Command(abc.ABC, ta.Generic[CommandInputT, CommandOutputT]):
-    @dc.dataclass(frozen=True)
-    class Input(abc.ABC):  # noqa
-        pass
-
+@dc.dataclass(frozen=True)
+class Command(abc.ABC, ta.Generic[CommandOutputT]):
     @dc.dataclass(frozen=True)
     class Output(abc.ABC):  # noqa
         pass
 
+
+##
+
+
+class CommandExecutor(abc.ABC, ta.Generic[CommandT, CommandOutputT]):
     @abc.abstractmethod
-    def _execute(self, inp: CommandInputT) -> CommandOutputT:
+    def execute(self, i: CommandT) -> CommandOutputT:
         raise NotImplementedError
 
 
 ########################################
-# ../../../pyremote.py
+# ../../pyremote.py
 """
 Basically this: https://mitogen.networkgenomics.com/howitworks.html
 """
@@ -203,6 +206,8 @@ _PYREMOTE_BOOTSTRAP_SRC_FD = 101
 
 _PYREMOTE_BOOTSTRAP_CHILD_PID_VAR = '_OPYR_CHILD_PID'
 _PYREMOTE_BOOTSTRAP_ARGV0_VAR = '_OPYR_ARGV0'
+_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR = '_OPYR_CONTEXT_NAME'
+_PYREMOTE_BOOTSTRAP_SRC_FILE_VAR = '_OPYR_SRC_FILE'
 _PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR = '_OPYR_OPTIONS_JSON'
 
 _PYREMOTE_BOOTSTRAP_ACK0 = b'OPYR000\n'
@@ -245,11 +250,10 @@ def _pyremote_bootstrap_main(context_name: str) -> None:
         for f in [r0, w0, r1, w1]:
             os.close(f)
 
-        # Save child pid to close after relaunch
+        # Save vars
         os.environ[_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR] = str(cp)
-
-        # Save original argv0
         os.environ[_PYREMOTE_BOOTSTRAP_ARGV0_VAR] = sys.executable
+        os.environ[_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,)))
@@ -300,6 +304,7 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
 
         '_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR',
         '_PYREMOTE_BOOTSTRAP_ARGV0_VAR',
+        '_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR',
 
         '_PYREMOTE_BOOTSTRAP_ACK0',
         '_PYREMOTE_BOOTSTRAP_ACK1',
@@ -335,14 +340,15 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
 class PyremotePayloadRuntime:
     input: ta.BinaryIO
     output: ta.BinaryIO
+    context_name: str
     main_src: str
     options: PyremoteBootstrapOptions
     env_info: PyremoteEnvInfo
 
 
 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:
+    # If src file var is not present we need to do initial finalization
+    if _PYREMOTE_BOOTSTRAP_SRC_FILE_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')
@@ -366,11 +372,14 @@ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
             os.write(tfd, main_src.encode('utf-8'))
             os.close(tfd)
 
-            # Set json options var
+            # Set vars
+            os.environ[_PYREMOTE_BOOTSTRAP_SRC_FILE_VAR] = tfn
             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)
+            exe = os.environ[_PYREMOTE_BOOTSTRAP_ARGV0_VAR]
+            context_name = os.environ[_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR]
+            os.execl(exe, exe + (_PYREMOTE_BOOTSTRAP_PROC_TITLE_FMT % (context_name,)), tfn)
 
     else:
         # Load options json var
@@ -378,12 +387,15 @@ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
         options = PyremoteBootstrapOptions(**json.loads(options_json_str))
 
         # Read temp source file
-        with open(sys.orig_argv[1]) as sf:
+        with open(os.environ.pop(_PYREMOTE_BOOTSTRAP_SRC_FILE_VAR)) as sf:
             main_src = sf.read()
 
     # Restore original argv0
     sys.executable = os.environ.pop(_PYREMOTE_BOOTSTRAP_ARGV0_VAR)
 
+    # Grab context name
+    context_name = os.environ.pop(_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR)
+
     # Write third ack
     os.write(1, _PYREMOTE_BOOTSTRAP_ACK2)
 
@@ -406,6 +418,7 @@ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
     return PyremotePayloadRuntime(
         input=input,
         output=output,
+        context_name=context_name,
         main_src=main_src,
         options=options,
         env_info=env_info,
@@ -526,7 +539,7 @@ class PyremoteBootstrapDriver:
 
 
 ########################################
-# ../../../../omlish/lite/cached.py
+# ../../../omlish/lite/cached.py
 
 
 class _cached_nullary:  # noqa
@@ -550,8 +563,14 @@ def cached_nullary(fn):  # ta.Callable[..., T]) -> ta.Callable[..., T]:
     return _cached_nullary(fn)
 
 
+def static_init(fn: CallableT) -> CallableT:
+    fn = cached_nullary(fn)
+    fn()
+    return fn
+
+
 ########################################
-# ../../../../omlish/lite/check.py
+# ../../../omlish/lite/check.py
 
 
 def check_isinstance(v: ta.Any, spec: ta.Union[ta.Type[T], tuple]) -> T:
@@ -648,7 +667,7 @@ def check_non_empty(v: SizedT) -> SizedT:
 
 
 ########################################
-# ../../../../omlish/lite/json.py
+# ../../../omlish/lite/json.py
 
 
 ##
@@ -679,7 +698,7 @@ json_dumps_compact: ta.Callable[..., str] = functools.partial(json.dumps, **JSON
 
 
 ########################################
-# ../../../../omlish/lite/reflect.py
+# ../../../omlish/lite/reflect.py
 
 
 _GENERIC_ALIAS_TYPES = (
@@ -734,7 +753,40 @@ def deep_subclasses(cls: ta.Type[T]) -> ta.Iterator[ta.Type[T]]:
 
 
 ########################################
-# ../../../../omlish/lite/logs.py
+# ../payload.py
+
+
+@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())
+
+
+def get_payload_src(*, file: ta.Optional[str]) -> str:
+    if file is not None:
+        with open(file) as f:
+            return f.read()
+
+    if _is_self_amalg():
+        return _get_self_src()
+
+    import importlib.resources
+    return importlib.resources.files(__package__.split('.')[0] + '.scripts').joinpath('manage.py').read_text()
+
+
+########################################
+# ../../../omlish/lite/logs.py
 """
 TODO:
  - translate json keys
@@ -1004,7 +1056,7 @@ def configure_standard_logging(
 
 
 ########################################
-# ../../../../omlish/lite/marshal.py
+# ../../../omlish/lite/marshal.py
 """
 TODO:
  - pickle stdlib objs? have to pin to 3.8 pickle protocol, will be cross-version
@@ -1351,7 +1403,7 @@ def unmarshal_obj(o: ta.Any, ty: ta.Union[ta.Type[T], ta.Any]) -> T:
 
 
 ########################################
-# ../../../../omlish/lite/runtime.py
+# ../../../omlish/lite/runtime.py
 
 
 @cached_nullary
@@ -1368,7 +1420,7 @@ def check_runtime_version() -> None:
 
 
 ########################################
-# ../../../../omlish/lite/subprocesses.py
+# ../../../omlish/lite/subprocesses.py
 
 
 ##
@@ -1498,24 +1550,23 @@ def subprocess_close(
 ##
 
 
-class SubprocessCommand(Command['SubprocessCommand.Input', 'SubprocessCommand.Output']):
-    @dc.dataclass(frozen=True)
-    class Input(Command.Input):
-        args: ta.Sequence[str]
+@dc.dataclass(frozen=True)
+class SubprocessCommand(Command['SubprocessCommand.Output']):
+    args: ta.Sequence[str]
 
-        shell: bool = False
-        cwd: ta.Optional[str] = None
-        env: ta.Optional[ta.Mapping[str, str]] = None
+    shell: bool = False
+    cwd: ta.Optional[str] = None
+    env: ta.Optional[ta.Mapping[str, str]] = None
 
-        capture_stdout: bool = False
-        capture_stderr: bool = False
+    capture_stdout: bool = False
+    capture_stderr: bool = False
 
-        input: ta.Optional[bytes] = None
-        timeout: ta.Optional[float] = None
+    input: ta.Optional[bytes] = None
+    timeout: ta.Optional[float] = None
 
-        def __post_init__(self) -> None:
-            if isinstance(self.args, str):
-                raise TypeError(self.args)
+    def __post_init__(self) -> None:
+        if isinstance(self.args, str):
+            raise TypeError(self.args)
 
     @dc.dataclass(frozen=True)
     class Output(Command.Output):
@@ -1527,7 +1578,12 @@ class SubprocessCommand(Command['SubprocessCommand.Input', 'SubprocessCommand.Ou
         stdout: ta.Optional[bytes] = None
         stderr: ta.Optional[bytes] = None
 
-    def _execute(self, inp: Input) -> Output:
+
+##
+
+
+class SubprocessCommandExecutor(CommandExecutor[SubprocessCommand, SubprocessCommand.Output]):
+    def execute(self, inp: SubprocessCommand) -> SubprocessCommand.Output:
         proc = subprocess.Popen(
             subprocess_maybe_shell_wrap_exec(*inp.args),
 
@@ -1558,6 +1614,101 @@ class SubprocessCommand(Command['SubprocessCommand.Input', 'SubprocessCommand.Ou
         )
 
 
+########################################
+# ../spawning.py
+
+
+class PySpawner:
+    DEFAULT_PYTHON = 'python3'
+
+    def __init__(
+            self,
+            src: str,
+            *,
+            shell: ta.Optional[str] = None,
+            shell_quote: bool = False,
+            python: str = DEFAULT_PYTHON,
+            stderr: ta.Optional[ta.Literal['pipe', 'stdout', 'devnull']] = None,
+    ) -> None:
+        super().__init__()
+
+        self._src = src
+        self._shell = shell
+        self._shell_quote = shell_quote
+        self._python = python
+        self._stderr = stderr
+
+    #
+
+    class _PreparedCmd(ta.NamedTuple):
+        cmd: ta.Sequence[str]
+        shell: bool
+
+    def _prepare_cmd(self) -> _PreparedCmd:
+        if self._shell is not None:
+            sh_src = f'{self._python} -c {shlex.quote(self._src)}'
+            if self._shell_quote:
+                sh_src = shlex.quote(sh_src)
+            sh_cmd = f'{self._shell} {sh_src}'
+            return PySpawner._PreparedCmd(
+                cmd=[sh_cmd],
+                shell=True,
+            )
+
+        else:
+            return PySpawner._PreparedCmd(
+                cmd=[self._python, '-c', self._src],
+                shell=False,
+            )
+
+    #
+
+    _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
+        stdout: ta.IO
+        stderr: ta.Optional[ta.IO]
+
+    @contextlib.contextmanager
+    def spawn(
+            self,
+            *,
+            timeout: ta.Optional[float] = None,
+    ) -> ta.Generator[Spawned, None, None]:
+        pc = self._prepare_cmd()
+
+        with subprocess.Popen(
+            subprocess_maybe_shell_wrap_exec(*pc.cmd),
+            shell=pc.shell,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=self._STDERR_KWARG_MAP[self._stderr] if self._stderr is not None else None,
+        ) as proc:
+            stdin = check_not_none(proc.stdin)
+            stdout = check_not_none(proc.stdout)
+
+            try:
+                yield PySpawner.Spawned(
+                    stdin=stdin,
+                    stdout=stdout,
+                    stderr=proc.stderr,
+                )
+
+            finally:
+                try:
+                    stdin.close()
+                except BrokenPipeError:
+                    pass
+
+                proc.wait(timeout)
+
+
 ########################################
 # main.py
 
@@ -1570,29 +1721,23 @@ _COMMAND_TYPES = {
 }
 
 
-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),
+@static_init
+def _register_command_marshaling() -> None:
+    for fn in [
+        lambda c: c,
+        lambda c: c.Output,
+    ]:
+        register_opj_marshaler(
+            fn(Command),
+            PolymorphicObjMarshaler.of([
+                PolymorphicObjMarshaler.Impl(
+                    fn(cty),
+                    k,
+                    get_obj_marshaler(fn(cty)),
+                )
+                for k, cty in _COMMAND_TYPES.items()
+            ]),
         )
-        for k, cty in _COMMAND_TYPES.items()
-    ]),
-)
 
 
 ##
@@ -1630,12 +1775,12 @@ def _remote_main() -> None:
     rt = pyremote_bootstrap_finalize()  # noqa
 
     while True:
-        i = _recv_obj(rt.input, Command.Input)
+        i = _recv_obj(rt.input, Command)
         if i is None:
             break
 
-        if isinstance(i, SubprocessCommand.Input):
-            o = SubprocessCommand()._execute(i)  # noqa
+        if isinstance(i, SubprocessCommand):
+            o = SubprocessCommandExecutor().execute(i)  # noqa
         else:
             raise TypeError(i)
 
@@ -1645,38 +1790,6 @@ def _remote_main() -> None:
 ##
 
 
-@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())
-
-
-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:
     import argparse
 
@@ -1685,19 +1798,20 @@ def _main() -> None:
     parser.add_argument('-s', '--shell')
     parser.add_argument('-q', '--shell-quote', action='store_true')
     parser.add_argument('--python', default='python3')
-    parser.add_argument('--_amalg-file')
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--_payload-file')
 
     args = parser.parse_args()
 
     #
 
-    amalg_src = _get_amalg_src(amalg_file=args._amalg_file)  # noqa
+    payload_src = get_payload_src(file=args._payload_file)  # noqa
 
     #
 
     remote_src = '\n\n'.join([
         '__name__ = "__remote__"',
-        amalg_src,
+        payload_src,
         '_remote_main()',
     ])
 
@@ -1705,60 +1819,41 @@ def _main() -> None:
 
     bs_src = pyremote_build_bootstrap_cmd(__package__ or 'manage')
 
-    if args.shell is not None:
-        sh_src = f'{args.python} -c {shlex.quote(bs_src)}'
-        if args.shell_quote:
-            sh_src = shlex.quote(sh_src)
-        sh_cmd = f'{args.shell} {sh_src}'
-        cmd = [sh_cmd]
-        shell = True
-    else:
-        cmd = [args.python, '-c', bs_src]
-        shell = False
-
-    proc = subprocess.Popen(
-        subprocess_maybe_shell_wrap_exec(*cmd),
-        shell=shell,
-        stdin=subprocess.PIPE,
-        stdout=subprocess.PIPE,
-    )
-
-    stdin = check_not_none(proc.stdin)
-    stdout = check_not_none(proc.stdout)
-
-    res = PyremoteBootstrapDriver(  # noqa
-        remote_src,
-        PyremoteBootstrapOptions(
-            # debug=True,
-        ),
-    ).run(stdin, stdout)
-    # print(res)
-
     #
 
-    for ci in [
-        SubprocessCommand.Input(
-            args=['python3', '-'],
-            input=b'print(1)\n',
-            capture_stdout=True,
-        ),
-        SubprocessCommand.Input(
-            args=['uname'],
-            capture_stdout=True,
-        ),
-    ]:
-        _send_obj(stdin, ci, Command.Input)
-
-        o = _recv_obj(stdout, Command.Output)
-
-        print(o)
+    spawner = PySpawner(
+        bs_src,
+        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)
 
-    try:
-        stdin.close()
-    except BrokenPipeError:
-        pass
+        #
 
-    proc.wait()
+        for ci in [
+            SubprocessCommand(
+                args=['python3', '-'],
+                input=b'print(1)\n',
+                capture_stdout=True,
+            ),
+            SubprocessCommand(
+                args=['uname'],
+                capture_stdout=True,
+            ),
+        ]:
+            _send_obj(proc.stdin, ci, Command)
+
+            o = _recv_obj(proc.stdout, Command.Output)
+
+            print(o)
 
 
 if __name__ == '__main__':

ominfra/scripts/supervisor.py

@@ -101,6 +101,7 @@ V = ta.TypeVar('V')
 
 # ../../omlish/lite/cached.py
 T = ta.TypeVar('T')
+CallableT = ta.TypeVar('CallableT', bound=ta.Callable)
 
 # ../../omlish/lite/check.py
 SizedT = ta.TypeVar('SizedT', bound=ta.Sized)
@@ -1454,6 +1455,12 @@ def cached_nullary(fn):  # ta.Callable[..., T]) -> ta.Callable[..., T]:
     return _cached_nullary(fn)
 
 
+def static_init(fn: CallableT) -> CallableT:
+    fn = cached_nullary(fn)
+    fn()
+    return fn
+
+
 ########################################
 # ../../../omlish/lite/check.py
 

{ominfra-0.0.0.dev138.dist-info → ominfra-0.0.0.dev140.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ominfra
-Version: 0.0.0.dev138
+Version: 0.0.0.dev140
 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.dev138
-Requires-Dist: omlish==0.0.0.dev138
+Requires-Dist: omdev==0.0.0.dev140
+Requires-Dist: omlish==0.0.0.dev140
 Provides-Extra: all
 Requires-Dist: paramiko~=3.5; extra == "all"
 Requires-Dist: asyncssh~=2.18; extra == "all"