ominfra 0.0.0.dev147__py3-none-any.whl → 0.0.0.dev148__py3-none-any.whl

ominfra/manage/remote/spawning.py

@@ -1,17 +1,23 @@
 # ruff: noqa: UP006 UP007
+import abc
+import asyncio
 import contextlib
 import dataclasses as dc
 import shlex
 import subprocess
 import typing as ta
 
+from omlish.lite.asyncio.subprocesses import asyncio_subprocess_popen
 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 RemoteSpawning:
+##
+
+
+class RemoteSpawning(abc.ABC):
     @dc.dataclass(frozen=True)
     class Target:
         shell: ta.Optional[str] = None
@@ -22,15 +28,35 @@ class RemoteSpawning:
 
         stderr: ta.Optional[str] = None  # SubprocessChannelOption
 
-    #
+    @dc.dataclass(frozen=True)
+    class Spawned:
+        stdin: asyncio.StreamWriter
+        stdout: asyncio.StreamReader
+        stderr: ta.Optional[asyncio.StreamReader]
+
+    @abc.abstractmethod
+    def spawn(
+            self,
+            tgt: Target,
+            src: str,
+            *,
+            timeout: ta.Optional[float] = None,
+            debug: bool = False,
+    ) -> ta.AsyncContextManager[Spawned]:
+        raise NotImplementedError
+
+
+##
 
-    class _PreparedCmd(ta.NamedTuple):
+
+class SubprocessRemoteSpawning(RemoteSpawning):
+    class _PreparedCmd(ta.NamedTuple):  # noqa
         cmd: ta.Sequence[str]
         shell: bool
 
     def _prepare_cmd(
             self,
-            tgt: Target,
+            tgt: RemoteSpawning.Target,
             src: str,
     ) -> _PreparedCmd:
         if tgt.shell is not None:
@@ -38,44 +64,38 @@ class RemoteSpawning:
             if tgt.shell_quote:
                 sh_src = shlex.quote(sh_src)
             sh_cmd = f'{tgt.shell} {sh_src}'
-            return RemoteSpawning._PreparedCmd(
-                cmd=[sh_cmd],
-                shell=True,
-            )
+            return SubprocessRemoteSpawning._PreparedCmd([sh_cmd], shell=True)
 
         else:
-            return RemoteSpawning._PreparedCmd(
-                cmd=[tgt.python, '-c', src],
-                shell=False,
-            )
+            return SubprocessRemoteSpawning._PreparedCmd([tgt.python, '-c', src], shell=False)
 
     #
 
-    @dc.dataclass(frozen=True)
-    class Spawned:
-        stdin: ta.IO
-        stdout: ta.IO
-        stderr: ta.Optional[ta.IO]
-
-    @contextlib.contextmanager
-    def spawn(
+    @contextlib.asynccontextmanager
+    async def spawn(
             self,
-            tgt: Target,
+            tgt: RemoteSpawning.Target,
             src: str,
             *,
             timeout: ta.Optional[float] = None,
-    ) -> ta.Generator[Spawned, None, None]:
+            debug: bool = False,
+    ) -> ta.AsyncGenerator[RemoteSpawning.Spawned, None]:
         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=(
-                SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, tgt.stderr)]
-                if tgt.stderr is not None else None
-            ),
+        cmd = pc.cmd
+        if not debug:
+            cmd = subprocess_maybe_shell_wrap_exec(*cmd)
+
+        async with asyncio_subprocess_popen(
+                *cmd,
+                shell=pc.shell,
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=(
+                    SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, tgt.stderr)]
+                    if tgt.stderr is not None else None
+                ),
+                timeout=timeout,
         ) as proc:
             stdin = check_not_none(proc.stdin)
             stdout = check_not_none(proc.stdout)
@@ -92,5 +112,3 @@ class RemoteSpawning:
                     stdin.close()
                 except BrokenPipeError:
                     pass
-
-                proc.wait(timeout)

ominfra/pyremote.py

@@ -324,10 +324,8 @@ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
         with open(os.environ.pop(_PYREMOTE_BOOTSTRAP_SRC_FILE_VAR)) as sf:
             main_src = sf.read()
 
-    # Restore original argv0
+    # Restore vars
     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
@@ -501,3 +499,30 @@ class PyremoteBootstrapDriver:
                 output.flush()
             else:
                 raise TypeError(go)
+
+    async def async_run(
+            self,
+            input: ta.Any,  # asyncio.StreamWriter  # noqa
+            output: ta.Any,  # asyncio.StreamReader
+    ) -> Result:
+        gen = self.gen()
+
+        gi: ta.Optional[bytes] = 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):
+                if len(gi := await input.read(go.sz)) != go.sz:
+                    raise EOFError
+            elif isinstance(go, self.Write):
+                gi = None
+                output.write(go.d)
+                await output.drain()
+            else:
+                raise TypeError(go)

ominfra/scripts/journald2aws.py

@@ -898,7 +898,10 @@ def toml_make_safe_parse_float(parse_float: TomlParseFloat) -> TomlParseFloat:
 # ../../../../../omlish/lite/cached.py
 
 
-class _cached_nullary:  # noqa
+##
+
+
+class _AbstractCachedNullary:
     def __init__(self, fn):
         super().__init__()
         self._fn = fn
@@ -906,17 +909,25 @@ class _cached_nullary:  # noqa
         functools.update_wrapper(self, fn)
 
     def __call__(self, *args, **kwargs):  # noqa
-        if self._value is self._missing:
-            self._value = self._fn()
-        return self._value
+        raise TypeError
 
     def __get__(self, instance, owner):  # noqa
         bound = instance.__dict__[self._fn.__name__] = self.__class__(self._fn.__get__(instance, owner))
         return bound
 
 
+##
+
+
+class _CachedNullary(_AbstractCachedNullary):
+    def __call__(self, *args, **kwargs):  # noqa
+        if self._value is self._missing:
+            self._value = self._fn()
+        return self._value
+
+
 def cached_nullary(fn):  # ta.Callable[..., T]) -> ta.Callable[..., T]:
-    return _cached_nullary(fn)
+    return _CachedNullary(fn)
 
 
 def static_init(fn: CallableT) -> CallableT:
@@ -925,6 +936,20 @@ def static_init(fn: CallableT) -> CallableT:
     return fn
 
 
+##
+
+
+class _AsyncCachedNullary(_AbstractCachedNullary):
+    async def __call__(self, *args, **kwargs):
+        if self._value is self._missing:
+            self._value = await self._fn()
+        return self._value
+
+
+def async_cached_nullary(fn):  # ta.Callable[..., T]) -> ta.Callable[..., T]:
+    return _AsyncCachedNullary(fn)
+
+
 ########################################
 # ../../../../../omlish/lite/check.py
 
@@ -2831,6 +2856,33 @@ class AwsLogMessageBuilder:
 # ../../../../configs.py
 
 
+def parse_config_file(
+        name: str,
+        f: ta.TextIO,
+) -> ConfigMapping:
+    if name.endswith('.toml'):
+        return toml_loads(f.read())
+
+    elif any(name.endswith(e) for e in ('.yml', '.yaml')):
+        yaml = __import__('yaml')
+        return yaml.safe_load(f)
+
+    elif name.endswith('.ini'):
+        import configparser
+        cp = configparser.ConfigParser()
+        cp.read_file(f)
+        config_dct: ta.Dict[str, ta.Any] = {}
+        for sec in cp.sections():
+            cd = config_dct
+            for k in sec.split('.'):
+                cd = cd.setdefault(k, {})
+            cd.update(cp.items(sec))
+        return config_dct
+
+    else:
+        return json.loads(f.read())
+
+
 def read_config_file(
         path: str,
         cls: ta.Type[T],
@@ -2838,13 +2890,10 @@ def read_config_file(
         prepare: ta.Optional[ta.Callable[[ConfigMapping], ConfigMapping]] = None,
 ) -> T:
     with open(path) as cf:
-        if path.endswith('.toml'):
-            config_dct = toml_loads(cf.read())
-        else:
-            config_dct = json.loads(cf.read())
+        config_dct = parse_config_file(os.path.basename(path), cf)
 
     if prepare is not None:
-        config_dct = prepare(config_dct)  # type: ignore
+        config_dct = prepare(config_dct)
 
     return unmarshal_obj(config_dct, cls)
 
@@ -3177,7 +3226,7 @@ def subprocess_maybe_shell_wrap_exec(*args: str) -> ta.Tuple[str, ...]:
         return args
 
 
-def _prepare_subprocess_invocation(
+def prepare_subprocess_invocation(
         *args: str,
         env: ta.Optional[ta.Mapping[str, ta.Any]] = None,
         extra_env: ta.Optional[ta.Mapping[str, ta.Any]] = None,
@@ -3185,9 +3234,9 @@ def _prepare_subprocess_invocation(
         shell: bool = False,
         **kwargs: ta.Any,
 ) -> ta.Tuple[ta.Tuple[ta.Any, ...], ta.Dict[str, ta.Any]]:
-    log.debug(args)
+    log.debug('prepare_subprocess_invocation: args=%r', args)
     if extra_env:
-        log.debug(extra_env)
+        log.debug('prepare_subprocess_invocation: extra_env=%r', extra_env)
 
     if extra_env:
         env = {**(env if env is not None else os.environ), **extra_env}
@@ -3206,14 +3255,46 @@ def _prepare_subprocess_invocation(
     )
 
 
-def subprocess_check_call(*args: str, stdout=sys.stderr, **kwargs: ta.Any) -> None:
-    args, kwargs = _prepare_subprocess_invocation(*args, stdout=stdout, **kwargs)
-    return subprocess.check_call(args, **kwargs)  # type: ignore
+##
 
 
-def subprocess_check_output(*args: str, **kwargs: ta.Any) -> bytes:
-    args, kwargs = _prepare_subprocess_invocation(*args, **kwargs)
-    return subprocess.check_output(args, **kwargs)
+@contextlib.contextmanager
+def subprocess_common_context(*args: ta.Any, **kwargs: ta.Any) -> ta.Iterator[None]:
+    start_time = time.time()
+    try:
+        log.debug('subprocess_common_context.try: args=%r', args)
+        yield
+
+    except Exception as exc:  # noqa
+        log.debug('subprocess_common_context.except: exc=%r', exc)
+        raise
+
+    finally:
+        end_time = time.time()
+        elapsed_s = end_time - start_time
+        log.debug('subprocess_common_context.finally: elapsed_s=%f args=%r', elapsed_s, args)
+
+
+##
+
+
+def subprocess_check_call(
+        *args: str,
+        stdout: ta.Any = sys.stderr,
+        **kwargs: ta.Any,
+) -> None:
+    args, kwargs = prepare_subprocess_invocation(*args, stdout=stdout, **kwargs)
+    with subprocess_common_context(*args, **kwargs):
+        return subprocess.check_call(args, **kwargs)  # type: ignore
+
+
+def subprocess_check_output(
+        *args: str,
+        **kwargs: ta.Any,
+) -> bytes:
+    args, kwargs = prepare_subprocess_invocation(*args, **kwargs)
+    with subprocess_common_context(*args, **kwargs):
+        return subprocess.check_output(args, **kwargs)
 
 
 def subprocess_check_output_str(*args: str, **kwargs: ta.Any) -> str:
@@ -3229,16 +3310,31 @@ DEFAULT_SUBPROCESS_TRY_EXCEPTIONS: ta.Tuple[ta.Type[Exception], ...] = (
 )
 
 
-def subprocess_try_call(
-        *args: str,
+def _subprocess_try_run(
+        fn: ta.Callable[..., T],
+        *args: ta.Any,
         try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
         **kwargs: ta.Any,
-) -> bool:
+) -> ta.Union[T, Exception]:
     try:
-        subprocess_check_call(*args, **kwargs)
+        return fn(*args, **kwargs)
     except try_exceptions as e:  # noqa
         if log.isEnabledFor(logging.DEBUG):
             log.exception('command failed')
+        return e
+
+
+def subprocess_try_call(
+        *args: str,
+        try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
+        **kwargs: ta.Any,
+) -> bool:
+    if isinstance(_subprocess_try_run(
+            subprocess_check_call,
+            *args,
+            try_exceptions=try_exceptions,
+            **kwargs,
+    ), Exception):
         return False
     else:
         return True
@@ -3249,12 +3345,15 @@ def subprocess_try_output(
         try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
         **kwargs: ta.Any,
 ) -> ta.Optional[bytes]:
-    try:
-        return subprocess_check_output(*args, **kwargs)
-    except try_exceptions as e:  # noqa
-        if log.isEnabledFor(logging.DEBUG):
-            log.exception('command failed')
+    if isinstance(ret := _subprocess_try_run(
+            subprocess_check_output,
+            *args,
+            try_exceptions=try_exceptions,
+            **kwargs,
+    ), Exception):
         return None
+    else:
+        return ret
 
 
 def subprocess_try_output_str(*args: str, **kwargs: ta.Any) -> ta.Optional[str]: