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

ominfra/manage/main.py

@@ -5,101 +5,18 @@
 manage.py -s 'docker run -i python:3.12'
 manage.py -s 'ssh -i /foo/bar.pem foo@bar.baz' -q --python=python3.8
 """
-import json
-import struct
-import typing as ta
-
-from omlish.lite.cached import static_init
-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 ..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
-from .commands.subprocess import SubprocessCommandExecutor
-from .payload import get_payload_src
-from .spawning import PySpawner
-
-
-##
-
-
-_COMMAND_TYPES = {
-    'subprocess': SubprocessCommand,
-}
-
-
-@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()
-            ]),
-        )
-
-
-##
-
-
-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)
-
-
-##
+from omlish.lite.logs import log  # noqa
+from omlish.lite.marshal import ObjMarshalOptions
+from omlish.lite.marshal import ObjMarshalerManager
+from omlish.lite.pycharm import PycharmRemoteDebug
 
-
-def _remote_main() -> None:
-    rt = pyremote_bootstrap_finalize()  # noqa
-
-    while True:
-        i = _recv_obj(rt.input, Command)
-        if i is None:
-            break
-
-        if isinstance(i, SubprocessCommand):
-            o = SubprocessCommandExecutor().execute(i)  # noqa
-        else:
-            raise TypeError(i)
-
-        _send_obj(rt.output, o, Command.Output)
+from .bootstrap import MainBootstrap
+from .bootstrap_ import main_bootstrap
+from .commands.subprocess import SubprocessCommand
+from .config import MainConfig
+from .remote.config import RemoteConfig
+from .remote.execution import RemoteExecution
+from .remote.spawning import RemoteSpawning
 
 
 ##
@@ -110,65 +27,70 @@ 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('--pycharm-debug-port', type=int)
+    parser.add_argument('--pycharm-debug-host')
+    parser.add_argument('--pycharm-debug-version')
+
     parser.add_argument('--debug', action='store_true')
-    parser.add_argument('--_payload-file')
+
+    parser.add_argument('command', nargs='+')
 
     args = parser.parse_args()
 
     #
 
-    payload_src = get_payload_src(file=args._payload_file)  # noqa
+    bs = MainBootstrap(
+        main_config=MainConfig(
+            log_level='DEBUG' if args.debug else 'INFO',
 
-    #
+            debug=bool(args.debug),
+        ),
+
+        remote_config=RemoteConfig(
+            payload_file=args._payload_file,  # noqa
 
-    remote_src = '\n\n'.join([
-        '__name__ = "__remote__"',
-        payload_src,
-        '_remote_main()',
-    ])
+            pycharm_remote_debug=PycharmRemoteDebug(
+                port=args.pycharm_debug_port,
+                host=args.pycharm_debug_host,
+                install_version=args.pycharm_debug_version,
+            ) if args.pycharm_debug_port is not None else None,
+        ),
+    )
+
+    injector = main_bootstrap(
+        bs,
+    )
 
     #
 
-    bs_src = pyremote_build_bootstrap_cmd(__package__ or 'manage')
+    cmds = [
+        SubprocessCommand([c])
+        for c in args.command
+    ]
 
     #
 
-    spawner = PySpawner(
-        bs_src,
+    tgt = RemoteSpawning.Target(
         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)
-
-        #
-
-        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)
+
+    with injector[RemoteExecution].connect(tgt, bs) as rce:
+        for cmd in cmds:
+            r = rce.try_execute(cmd)
+
+            print(injector[ObjMarshalerManager].marshal_obj(r, opts=ObjMarshalOptions(raw_bytes=True)))
+
+    #
+
+    print('Success')
 
 
 if __name__ == '__main__':

ominfra/manage/marshal.py

@@ -0,0 +1,12 @@
+import dataclasses as dc
+import typing as ta
+
+from omlish.lite.marshal import ObjMarshalerManager
+
+
+@dc.dataclass(frozen=True)
+class ObjMarshalerInstaller:
+    fn: ta.Callable[[ObjMarshalerManager], None]
+
+
+ObjMarshalerInstallers = ta.NewType('ObjMarshalerInstallers', ta.Sequence[ObjMarshalerInstaller])

ominfra/manage/remote/channel.py

@@ -0,0 +1,52 @@
+# ruff: noqa: UP006 UP007
+import json
+import struct
+import typing as ta
+
+from omlish.lite.json import json_dumps_compact
+from omlish.lite.marshal import OBJ_MARSHALER_MANAGER
+from omlish.lite.marshal import ObjMarshalerManager
+
+
+T = ta.TypeVar('T')
+
+
+class RemoteChannel:
+    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 set_marshaler(self, msh: ObjMarshalerManager) -> None:
+        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.Type[T]) -> ta.Optional[T]:
+        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)

ominfra/manage/remote/config.py

@@ -0,0 +1,12 @@
+# ruff: noqa: UP006 UP007
+import dataclasses as dc
+import typing as ta
+
+from omlish.lite.pycharm import PycharmRemoteDebug
+
+
+@dc.dataclass(frozen=True)
+class RemoteConfig:
+    payload_file: ta.Optional[str] = None
+
+    pycharm_remote_debug: ta.Optional[PycharmRemoteDebug] = None

ominfra/manage/remote/execution.py

@@ -0,0 +1,193 @@
+# ruff: noqa: UP006 UP007
+import contextlib
+import dataclasses as dc
+import logging
+import typing as ta
+
+from omlish.lite.cached import cached_nullary
+from omlish.lite.check import check_not_none
+from omlish.lite.logs import log
+from omlish.lite.marshal import ObjMarshalerManager
+from omlish.lite.pycharm import pycharm_debug_connect
+
+from ...pyremote import PyremoteBootstrapDriver
+from ...pyremote import PyremoteBootstrapOptions
+from ...pyremote import pyremote_bootstrap_finalize
+from ...pyremote import pyremote_build_bootstrap_cmd
+from ..bootstrap import MainBootstrap
+from ..commands.base import Command
+from ..commands.base import CommandException
+from ..commands.base import CommandExecutor
+from ..commands.base import CommandOutputOrException
+from ..commands.base import CommandOutputOrExceptionData
+from .channel import RemoteChannel
+from .payload import RemoteExecutionPayloadFile
+from .payload import get_remote_payload_src
+from .spawning import RemoteSpawning
+
+
+if ta.TYPE_CHECKING:
+    from ..bootstrap_ import main_bootstrap
+else:
+    main_bootstrap: ta.Any = None
+
+
+##
+
+
+def _remote_execution_main() -> None:
+    rt = pyremote_bootstrap_finalize()  # noqa
+
+    chan = RemoteChannel(
+        rt.input,
+        rt.output,
+    )
+
+    bs = check_not_none(chan.recv_obj(MainBootstrap))
+
+    if (prd := bs.remote_config.pycharm_remote_debug) is not None:
+        pycharm_debug_connect(prd)
+
+    injector = main_bootstrap(bs)
+
+    chan.set_marshaler(injector[ObjMarshalerManager])
+
+    ce = injector[CommandExecutor]
+
+    while True:
+        i = chan.recv_obj(Command)
+        if i is None:
+            break
+
+        r = ce.try_execute(
+            i,
+            log=log,
+            omit_exc_object=True,
+        )
+
+        chan.send_obj(r)
+
+
+##
+
+
+@dc.dataclass()
+class RemoteCommandError(Exception):
+    e: CommandException
+
+
+class RemoteCommandExecutor(CommandExecutor):
+    def __init__(self, chan: RemoteChannel) -> None:
+        super().__init__()
+
+        self._chan = chan
+
+    def _remote_execute(self, cmd: Command) -> CommandOutputOrException:
+        self._chan.send_obj(cmd, Command)
+
+        if (r := self._chan.recv_obj(CommandOutputOrExceptionData)) is None:
+            raise EOFError
+
+        return r
+
+    # @ta.override
+    def execute(self, cmd: Command) -> Command.Output:
+        r = self._remote_execute(cmd)
+        if (e := r.exception) is not None:
+            raise RemoteCommandError(e)
+        else:
+            return check_not_none(r.output)
+
+    # @ta.override
+    def try_execute(
+            self,
+            cmd: Command,
+            *,
+            log: ta.Optional[logging.Logger] = None,
+            omit_exc_object: bool = False,
+    ) -> CommandOutputOrException:
+        try:
+            r = self._remote_execute(cmd)
+
+        except Exception as e:  # noqa
+            if log is not None:
+                log.exception('Exception executing remote command: %r', type(cmd))
+
+            return CommandOutputOrExceptionData(exception=CommandException.of(
+                e,
+                omit_exc_object=omit_exc_object,
+                cmd=cmd,
+            ))
+
+        else:
+            return r
+
+
+##
+
+
+class RemoteExecution:
+    def __init__(
+            self,
+            *,
+            spawning: RemoteSpawning,
+            msh: ObjMarshalerManager,
+            payload_file: ta.Optional[RemoteExecutionPayloadFile] = None,
+    ) -> None:
+        super().__init__()
+
+        self._spawning = spawning
+        self._msh = msh
+        self._payload_file = payload_file
+
+    #
+
+    @cached_nullary
+    def _payload_src(self) -> str:
+        return get_remote_payload_src(file=self._payload_file)
+
+    @cached_nullary
+    def _remote_src(self) -> ta.Sequence[str]:
+        return [
+            self._payload_src(),
+            '_remote_execution_main()',
+        ]
+
+    @cached_nullary
+    def _spawn_src(self) -> str:
+        return pyremote_build_bootstrap_cmd(__package__ or 'manage')
+
+    #
+
+    @contextlib.contextmanager
+    def connect(
+            self,
+            tgt: RemoteSpawning.Target,
+            bs: MainBootstrap,
+    ) -> ta.Generator[RemoteCommandExecutor, None, None]:
+        spawn_src = self._spawn_src()
+        remote_src = self._remote_src()
+
+        with self._spawning.spawn(
+                tgt,
+                spawn_src,
+        ) as proc:
+            res = PyremoteBootstrapDriver(  # noqa
+                remote_src,
+                PyremoteBootstrapOptions(
+                    debug=bs.main_config.debug,
+                ),
+            ).run(
+                proc.stdout,
+                proc.stdin,
+            )
+
+            chan = RemoteChannel(
+                proc.stdout,
+                proc.stdin,
+                msh=self._msh,
+            )
+
+            chan.send_obj(bs)
+
+            yield RemoteCommandExecutor(chan)

ominfra/manage/remote/inject.py

@@ -0,0 +1,29 @@
+# ruff: noqa: UP006 UP007
+import typing as ta
+
+from omlish.lite.inject import InjectorBindingOrBindings
+from omlish.lite.inject import InjectorBindings
+from omlish.lite.inject import inj
+
+from .config import RemoteConfig
+from .execution import RemoteExecution
+from .payload import RemoteExecutionPayloadFile
+from .spawning import RemoteSpawning
+
+
+def bind_remote(
+        *,
+        remote_config: RemoteConfig,
+) -> InjectorBindings:
+    lst: ta.List[InjectorBindingOrBindings] = [
+        inj.bind(remote_config),
+
+        inj.bind(RemoteSpawning, singleton=True),
+
+        inj.bind(RemoteExecution, singleton=True),
+    ]
+
+    if (pf := remote_config.payload_file) is not None:
+        lst.append(inj.bind(pf, to_key=RemoteExecutionPayloadFile))
+
+    return inj.as_bindings(*lst)

ominfra/manage/{payload.py → remote/payload.py}

@@ -6,6 +6,9 @@ import typing as ta
 from omlish.lite.cached import cached_nullary
 
 
+RemoteExecutionPayloadFile = ta.NewType('RemoteExecutionPayloadFile', str)
+
+
 @cached_nullary
 def _get_self_src() -> str:
     return inspect.getsource(sys.modules[__name__])
@@ -23,7 +26,10 @@ def _is_self_amalg() -> bool:
     return _is_src_amalg(_get_self_src())
 
 
-def get_payload_src(*, file: ta.Optional[str]) -> str:
+def get_remote_payload_src(
+        *,
+        file: ta.Optional[RemoteExecutionPayloadFile],
+) -> str:
     if file is not None:
         with open(file) as f:
             return f.read()

ominfra/manage/{spawning.py → remote/spawning.py}

@@ -6,28 +6,21 @@ import subprocess
 import typing as ta
 
 from omlish.lite.check import check_not_none
+from omlish.lite.subprocesses import SUBPROCESS_CHANNEL_OPTION_VALUES
+from omlish.lite.subprocesses import SubprocessChannelOption
 from omlish.lite.subprocesses import subprocess_maybe_shell_wrap_exec
 
 
-class PySpawner:
-    DEFAULT_PYTHON = 'python3'
+class RemoteSpawning:
+    @dc.dataclass(frozen=True)
+    class Target:
+        shell: ta.Optional[str] = None
+        shell_quote: bool = False
 
-    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
+        DEFAULT_PYTHON: ta.ClassVar[str] = 'python3'
+        python: str = DEFAULT_PYTHON
+
+        stderr: ta.Optional[str] = None  # SubprocessChannelOption
 
     #
 
@@ -35,31 +28,29 @@ class PySpawner:
         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:
+    def _prepare_cmd(
+            self,
+            tgt: Target,
+            src: str,
+    ) -> _PreparedCmd:
+        if tgt.shell is not None:
+            sh_src = f'{tgt.python} -c {shlex.quote(src)}'
+            if tgt.shell_quote:
                 sh_src = shlex.quote(sh_src)
-            sh_cmd = f'{self._shell} {sh_src}'
-            return PySpawner._PreparedCmd(
+            sh_cmd = f'{tgt.shell} {sh_src}'
+            return RemoteSpawning._PreparedCmd(
                 cmd=[sh_cmd],
                 shell=True,
             )
 
         else:
-            return PySpawner._PreparedCmd(
-                cmd=[self._python, '-c', self._src],
+            return RemoteSpawning._PreparedCmd(
+                cmd=[tgt.python, '-c', 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
@@ -69,23 +60,28 @@ class PySpawner:
     @contextlib.contextmanager
     def spawn(
             self,
+            tgt: Target,
+            src: str,
             *,
             timeout: ta.Optional[float] = None,
     ) -> ta.Generator[Spawned, None, None]:
-        pc = self._prepare_cmd()
+        pc = self._prepare_cmd(tgt, src)
 
         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,
+            stderr=(
+                SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, tgt.stderr)]
+                if tgt.stderr is not None else None
+            ),
         ) as proc:
             stdin = check_not_none(proc.stdin)
             stdout = check_not_none(proc.stdout)
 
             try:
-                yield PySpawner.Spawned(
+                yield RemoteSpawning.Spawned(
                     stdin=stdin,
                     stdout=stdout,
                     stderr=proc.stderr,