ominfra 0.0.0.dev102__py3-none-any.whl → 0.0.0.dev104__py3-none-any.whl

ominfra/pyremote/_runcommands.py

@@ -41,7 +41,7 @@ if sys.version_info < (3, 8):
 ########################################
 
 
-# ../../../omlish/lite/check.py
+# ../../../omlish/lite/cached.py
 T = ta.TypeVar('T')
 
 
@@ -195,7 +195,7 @@ def post_boostrap() -> PostBoostrap:
 # ../../../omlish/lite/cached.py
 
 
-class cached_nullary:  # noqa
+class _cached_nullary:  # noqa
     def __init__(self, fn):
         super().__init__()
         self._fn = fn
@@ -212,6 +212,10 @@ class cached_nullary:  # noqa
         return bound
 
 
+def cached_nullary(fn: ta.Callable[..., T]) -> ta.Callable[..., T]:
+    return _cached_nullary(fn)
+
+
 ########################################
 # ../../../omlish/lite/check.py
 
@@ -744,7 +748,7 @@ class DataclassObjMarshaler(ObjMarshaler):
         return {k: m.marshal(getattr(o, k)) for k, m in self.fs.items()}
 
     def unmarshal(self, o: ta.Any) -> ta.Any:
-        return self.ty(**{k: self.fs[k].unmarshal(v) for k, v in o.items() if self.nonstrict or k in self.fs})
+        return self.ty(**{k: self.fs[k].unmarshal(v) for k, v in o.items() if not self.nonstrict or k in self.fs})
 
 
 @dc.dataclass(frozen=True)
@@ -804,7 +808,10 @@ class UuidObjMarshaler(ObjMarshaler):
         return uuid.UUID(o)
 
 
-_OBJ_MARSHALERS: ta.Dict[ta.Any, ObjMarshaler] = {
+##
+
+
+_DEFAULT_OBJ_MARSHALERS: ta.Dict[ta.Any, ObjMarshaler] = {
     **{t: NopObjMarshaler() for t in (type(None),)},
     **{t: CastObjMarshaler(t) for t in (int, float, str, bool)},
     **{t: Base64ObjMarshaler(t) for t in (bytes, bytearray)},
@@ -833,20 +840,19 @@ _OBJ_MARSHALER_GENERIC_ITERABLE_TYPES: ta.Dict[ta.Any, type] = {
 }
 
 
-def register_opj_marshaler(ty: ta.Any, m: ObjMarshaler) -> None:
-    if ty in _OBJ_MARSHALERS:
-        raise KeyError(ty)
-    _OBJ_MARSHALERS[ty] = m
-
-
-def _make_obj_marshaler(ty: ta.Any) -> ObjMarshaler:
+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__:
             impls = [  # type: ignore
                 PolymorphicObjMarshaler.Impl(
                     ity,
                     ity.__qualname__,
-                    get_obj_marshaler(ity),
+                    rec(ity),
                 )
                 for ity in deep_subclasses(ty)
                 if abc.ABC not in ity.__bases__
@@ -862,7 +868,8 @@ def _make_obj_marshaler(ty: ta.Any) -> ObjMarshaler:
         if dc.is_dataclass(ty):
             return DataclassObjMarshaler(
                 ty,
-                {f.name: get_obj_marshaler(f.type) for f in dc.fields(ty)},
+                {f.name: rec(f.type) for f in dc.fields(ty)},
+                nonstrict=nonstrict_dataclasses,
             )
 
     if is_generic_alias(ty):
@@ -872,7 +879,7 @@ def _make_obj_marshaler(ty: ta.Any) -> ObjMarshaler:
             pass
         else:
             k, v = ta.get_args(ty)
-            return MappingObjMarshaler(mt, get_obj_marshaler(k), get_obj_marshaler(v))
+            return MappingObjMarshaler(mt, rec(k), rec(v))
 
         try:
             st = _OBJ_MARSHALER_GENERIC_ITERABLE_TYPES[ta.get_origin(ty)]
@@ -880,33 +887,71 @@ def _make_obj_marshaler(ty: ta.Any) -> ObjMarshaler:
             pass
         else:
             [e] = ta.get_args(ty)
-            return IterableObjMarshaler(st, get_obj_marshaler(e))
+            return IterableObjMarshaler(st, rec(e))
 
         if is_union_alias(ty):
-            return OptionalObjMarshaler(get_obj_marshaler(get_optional_alias_arg(ty)))
+            return OptionalObjMarshaler(rec(get_optional_alias_arg(ty)))
 
     raise TypeError(ty)
 
 
-def get_obj_marshaler(ty: ta.Any) -> ObjMarshaler:
-    try:
-        return _OBJ_MARSHALERS[ty]
-    except KeyError:
-        pass
+##
 
-    p = ProxyObjMarshaler()
-    _OBJ_MARSHALERS[ty] = p
-    try:
-        m = _make_obj_marshaler(ty)
-    except Exception:
-        del _OBJ_MARSHALERS[ty]
-        raise
-    else:
-        p.m = m
+
+_OBJ_MARSHALERS_LOCK = threading.RLock()
+
+_OBJ_MARSHALERS: ta.Dict[ta.Any, ObjMarshaler] = dict(_DEFAULT_OBJ_MARSHALERS)
+
+_OBJ_MARSHALER_PROXIES: ta.Dict[ta.Any, ProxyObjMarshaler] = {}
+
+
+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 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]
+            except KeyError:
+                pass
+
+        try:
+            return _OBJ_MARSHALER_PROXIES[ty]
+        except KeyError:
+            pass
+
+        rec = functools.partial(
+            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
+
+        if not no_cache:
+            _OBJ_MARSHALERS[ty] = m
         return m
 
 
+##
+
+
 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)
 
@@ -1039,6 +1084,24 @@ def subprocess_try_output_str(*args: str, **kwargs: ta.Any) -> ta.Optional[str]:
     return out.decode().strip() if out is not None else None
 
 
+##
+
+
+def subprocess_close(
+        proc: subprocess.Popen,
+        timeout: ta.Optional[float] = None,
+) -> None:
+    # TODO: terminate, sleep, kill
+    if proc.stdout:
+        proc.stdout.close()
+    if proc.stderr:
+        proc.stderr.close()
+    if proc.stdin:
+        proc.stdin.close()
+
+    proc.wait(timeout)
+
+
 ########################################
 # runcommands.py