omlish 0.0.0.dev443__py3-none-any.whl → 0.0.0.dev445__py3-none-any.whl

omlish/lang/imports/capture.py

@@ -1,6 +1,20 @@
+"""
+Insufficient alt impls:
+ - `sys.meta_path`:
+  - need access to the `fromlist` and `level` arguments passed to `__import__`
+  - need to return fake modules to the import operation which are not added to `sys.modules`
+ - `sys.addaudithook`: cannot prevent import or inject result
+ - `sys.settrace` / bytecode tracing: same
+ - jit bytecode rewriting: slower than just importing everything
+
+Possible alt impls:
+ - aot static analysis, codegen, compare, if valid skip ctxmgr body and inject proxies, otherwise warn and run
+"""
 import builtins
 import contextlib
 import functools
+import sys
+import threading
 import types
 import typing as ta
 
@@ -58,7 +72,10 @@ class ImportCaptureErrors:
         pass
 
 
-class _ImportCaptureImpl:
+##
+
+
+class _ImportCaptureHook:
     class ModuleSpec(ta.NamedTuple):
         name: str
         level: int
@@ -69,20 +86,26 @@ class _ImportCaptureImpl:
         def __repr__(self) -> str:
             return repr(str(self))
 
-    def __init__(self) -> None:
+    def __init__(
+            self,
+            *,
+            forbid_uncaptured_imports: bool = False,
+    ) -> None:
         super().__init__()
 
-        self._modules_by_spec: dict[_ImportCaptureImpl.ModuleSpec, _ImportCaptureImpl._Module] = {}
-        self._modules_by_module_obj: dict[types.ModuleType, _ImportCaptureImpl._Module] = {}
+        self._forbid_uncaptured_imports = forbid_uncaptured_imports
+
+        self._modules_by_spec: dict[_ImportCaptureHook.ModuleSpec, _ImportCaptureHook._Module] = {}
+        self._modules_by_module_obj: dict[types.ModuleType, _ImportCaptureHook._Module] = {}
 
-        self._attrs: dict[_ImportCaptureImpl._ModuleAttr, tuple[_ImportCaptureImpl._Module, str]] = {}
+        self._attrs: dict[_ImportCaptureHook._ModuleAttr, tuple[_ImportCaptureHook._Module, str]] = {}
 
     #
 
     class _ModuleAttr:
         def __init__(
                 self,
-                module: '_ImportCaptureImpl._Module',
+                module: '_ImportCaptureHook._Module',
                 name: str,
         ) -> None:
             super().__init__()
@@ -96,9 +119,9 @@ class _ImportCaptureImpl:
     class _Module:
         def __init__(
                 self,
-                spec: '_ImportCaptureImpl.ModuleSpec',
+                spec: '_ImportCaptureHook.ModuleSpec',
                 *,
-                getattr_handler: ta.Callable[['_ImportCaptureImpl._Module', str], ta.Any] | None = None,
+                getattr_handler: ta.Callable[['_ImportCaptureHook._Module', str], ta.Any] | None = None,
         ) -> None:
             super().__init__()
 
@@ -109,7 +132,7 @@ class _ImportCaptureImpl:
                 self.module_obj.__getattr__ = functools.partial(getattr_handler, self)  # type: ignore[method-assign]  # noqa
             self.initial_module_dict = dict(self.module_obj.__dict__)
 
-            self.contents: dict[str, _ImportCaptureImpl._ModuleAttr | types.ModuleType] = {}
+            self.contents: dict[str, _ImportCaptureHook._ModuleAttr | types.ModuleType] = {}
             self.imported_whole = False
 
         def __repr__(self) -> str:
@@ -133,17 +156,17 @@ class _ImportCaptureImpl:
         if attr in module.contents:
             raise ImportCaptureErrors.AttrError(str(module.spec), attr)
 
-        v: _ImportCaptureImpl._ModuleAttr | types.ModuleType
+        v: _ImportCaptureHook._ModuleAttr | types.ModuleType
         if not module.spec.name:
             if not module.spec.level:
                 raise ImportCaptureError
-            cs = _ImportCaptureImpl.ModuleSpec(attr, module.spec.level)
+            cs = _ImportCaptureHook.ModuleSpec(attr, module.spec.level)
             cm = self._get_or_make_module(cs)
             cm.imported_whole = True
             v = cm.module_obj
 
         else:
-            ma = _ImportCaptureImpl._ModuleAttr(module, attr)
+            ma = _ImportCaptureHook._ModuleAttr(module, attr)
             self._attrs[ma] = (module, attr)
             v = ma
 
@@ -173,7 +196,7 @@ class _ImportCaptureImpl:
                 bad = False
                 if x is not module.contents.get(attr):
                     bad = True
-                if isinstance(x, _ImportCaptureImpl._ModuleAttr):
+                if isinstance(x, _ImportCaptureHook._ModuleAttr):
                     if self._attrs[x] != (module, attr):
                         bad = True
                 elif isinstance(x, types.ModuleType):
@@ -202,7 +225,7 @@ class _ImportCaptureImpl:
         ):
             return None
 
-        spec = _ImportCaptureImpl.ModuleSpec(name, level)
+        spec = _ImportCaptureHook.ModuleSpec(name, level)
         module = self._get_or_make_module(spec)
 
         self._handle_import(
@@ -212,64 +235,33 @@ class _ImportCaptureImpl:
 
         return module.module_obj
 
+    @ta.final
     @contextlib.contextmanager
     def hook_context(
             self,
             mod_globals: ta.MutableMapping[str, ta.Any],  # noqa
-            *,
-            forbid_uncaptured_imports: bool = False,
     ) -> ta.Iterator[None]:
         if self._MOD_SELF_ATTR in mod_globals:
             raise ImportCaptureErrors.HookError
 
-        old_import = builtins.__import__
-
-        def new_import(
-                name,
-                globals=None,  # noqa
-                locals=None,  # noqa
-                fromlist=None,
-                level=0,
-        ):
-            if (im := self._intercept_import(
-                    name,
-                    globals=globals,
-                    from_list=fromlist,
-                    level=level,
-            )) is not None:
-                return im
-
-            if forbid_uncaptured_imports:
-                raise ImportCaptureErrors.UncapturedImportForbiddenError(
-                    str(_ImportCaptureImpl.ModuleSpec(name, level)),
-                    fromlist,
-                )
-
-            return old_import(
-                name,
-                globals=globals,
-                locals=locals,
-                fromlist=fromlist,
-                level=level,
-            )
-
-        #
-
         mod_globals[self._MOD_SELF_ATTR] = self
-        builtins.__import__ = new_import
 
         try:
-            yield
+            with self._hook_context(mod_globals):
+                yield
 
         finally:
-            if not (
-                    mod_globals[self._MOD_SELF_ATTR] is self and
-                    builtins.__import__ is new_import
-            ):
+            if mod_globals[self._MOD_SELF_ATTR] is not self:
                 raise ImportCaptureErrors.HookError
 
             del mod_globals[self._MOD_SELF_ATTR]
-            builtins.__import__ = old_import
+
+    # @abc.abstractmethod
+    def _hook_context(
+            self,
+            mod_globals: ta.MutableMapping[str, ta.Any],  # noqa
+    ) -> ta.ContextManager[None]:
+        raise NotImplementedError
 
     #
 
@@ -298,16 +290,16 @@ class _ImportCaptureImpl:
             *,
             collect_unreferenced: bool = False,
     ) -> 'ImportCapture.Captured':
-        dct: dict[_ImportCaptureImpl._Module, list[tuple[str | None, str]]] = {}
+        dct: dict[_ImportCaptureHook._Module, list[tuple[str | None, str]]] = {}
 
-        rem_whole_mods: set[_ImportCaptureImpl._Module] = set()
-        rem_mod_attrs: set[_ImportCaptureImpl._ModuleAttr] = set()
+        rem_whole_mods: set[_ImportCaptureHook._Module] = set()
+        rem_mod_attrs: set[_ImportCaptureHook._ModuleAttr] = set()
         if collect_unreferenced:
             rem_whole_mods.update([m for m in self._modules_by_spec.values() if m.imported_whole])
             rem_mod_attrs.update(self._attrs)
 
         for attr, obj in mod_globals.items():
-            if isinstance(obj, _ImportCaptureImpl._ModuleAttr):
+            if isinstance(obj, _ImportCaptureHook._ModuleAttr):
                 try:
                     m, a = self._attrs[obj]
                 except KeyError:
@@ -315,7 +307,7 @@ class _ImportCaptureImpl:
                 dct.setdefault(m, []).append((a, attr))
                 rem_mod_attrs.discard(obj)
 
-            elif isinstance(obj, _ImportCaptureImpl._Module):
+            elif isinstance(obj, _ImportCaptureHook._Module):
                 raise ImportCaptureErrors.AttrError(None, attr) from None
 
             elif isinstance(obj, types.ModuleType):
@@ -362,14 +354,236 @@ class _ImportCaptureImpl:
         )
 
 
-class ImportCapture:
-    """
-    This is a bit extreme, but worth it. For simplicity, it currently relies on temporarily patching
-    `__builtins__.__import__` for the duration of its context manager, but it can be switched to use any number of other
-    import hooks (like `sys.meta_path`). It does not rely on any permanent modification to import machinery, only for
-    the duration of its capture.
-    """
+#
+
+
+class _AbstractBuiltinsImportCaptureHook(_ImportCaptureHook):
+    def _new_import(
+            self,
+            old_import,
+            name,
+            globals=None,  # noqa
+            locals=None,  # noqa
+            fromlist=None,
+            level=0,
+    ):
+        if (im := self._intercept_import(
+                name,
+                globals=globals,
+                from_list=fromlist,
+                level=level,
+        )) is not None:
+            return im
+
+        if self._forbid_uncaptured_imports:
+            raise ImportCaptureErrors.UncapturedImportForbiddenError(
+                str(_ImportCaptureHook.ModuleSpec(name, level)),
+                fromlist,
+            )
+
+        return old_import(
+            name,
+            globals=globals,
+            locals=locals,
+            fromlist=fromlist,
+            level=level,
+        )
+
+
+class _UnsafeGlobalBuiltinsImportCaptureHook(_AbstractBuiltinsImportCaptureHook):
+    @contextlib.contextmanager
+    def _hook_context(
+            self,
+            mod_globals: ta.MutableMapping[str, ta.Any],  # noqa
+    ) -> ta.Iterator[None]:
+        old_import = builtins.__import__
+        new_import = functools.partial(self._new_import, old_import)
+
+        builtins.__import__ = new_import
+
+        try:
+            yield
+
+        finally:
+            if builtins.__import__ is not new_import:
+                raise ImportCaptureErrors.HookError
+
+            builtins.__import__ = old_import
+
+
+class _SomewhatThreadSafeGlobalBuiltinsImportCaptureHook(_AbstractBuiltinsImportCaptureHook):
+    class _AlreadyPatchedError(Exception):
+        pass
+
+    @ta.final
+    class _Patch:
+        __lock: ta.ClassVar[threading.Lock] = threading.Lock()
+
+        def __init__(self, old_import):
+            self.__old_import = old_import
+            self.__hooks = {}
+            self.__uninstalled = False
+
+        @classmethod
+        def _add_hook(cls, mod_globals, new_import) -> '_SomewhatThreadSafeGlobalBuiltinsImportCaptureHook._Patch':
+            gi = id(mod_globals)
+            for _ in range(1_000):
+                try:
+                    with cls.__lock:
+                        x: ta.Any = builtins.__import__
+                        p: _SomewhatThreadSafeGlobalBuiltinsImportCaptureHook._Patch
+                        if x.__class__ is cls:
+                            p = x
+                            if p.__uninstalled:  # noqa
+                                raise _SomewhatThreadSafeGlobalBuiltinsImportCaptureHook._AlreadyPatchedError  # noqa
+                        else:
+                            p = cls(x)
+                            builtins.__import__ = p
+                        p.__hooks[gi] = (mod_globals, new_import)
+                        return p
+                except _SomewhatThreadSafeGlobalBuiltinsImportCaptureHook._AlreadyPatchedError:
+                    pass
+            raise ImportCaptureErrors.HookError('Failed to install builtins hook')
+
+        def _remove_hook(self, mod_globals, *, no_raise=False):
+            gi = id(mod_globals)
+            with self.__lock:
+                tg, _ = self.__hooks[gi]
+                del self.__hooks[gi]
+                if not self.__uninstalled and not self.__hooks:
+                    self.__uninstalled = True
+                    if builtins.__import__ is not self:
+                        if not no_raise:
+                            # TODO: warn?
+                            raise ImportCaptureErrors.HookError('Unexpected builtins hook')
+                    else:
+                        builtins.__import__ = self.__old_import
+                if tg is not mod_globals:
+                    if not no_raise:
+                        # TODO: warn?
+                        raise ImportCaptureErrors.HookError('Mismatched globals')
+
+        def __call__(
+                self,
+                name,
+                globals=None,  # noqa
+                locals=None,  # noqa
+                fromlist=None,
+                level=0,
+        ):
+            if globals is not None and (tup := self.__hooks.get(id(globals))) is not None:
+                tg, tf = tup
+                if tg is globals:
+                    return tf(
+                        self.__old_import,
+                        name,
+                        globals=globals,
+                        locals=locals,
+                        fromlist=fromlist,
+                        level=level,
+                    )
+                else:
+                    self._remove_hook(tg, no_raise=True)
+
+            return self.__old_import(
+                name,
+                globals=globals,
+                locals=locals,
+                fromlist=fromlist,
+                level=level,
+            )
+
+    @contextlib.contextmanager
+    def _hook_context(
+            self,
+            mod_globals: ta.MutableMapping[str, ta.Any],  # noqa
+    ) -> ta.Iterator[None]:
+        patch = _SomewhatThreadSafeGlobalBuiltinsImportCaptureHook._Patch._add_hook(mod_globals, self._new_import)  # noqa
+
+        try:
+            yield
+
+        finally:
+            patch._remove_hook(mod_globals)  # noqa
 
+
+#
+
+
+_capture: ta.Any = None
+try:
+    from . import _capture  # type: ignore
+except ImportError:
+    pass
+
+
+class _FrameBuiltinsImportCaptureHook(_AbstractBuiltinsImportCaptureHook):
+    def __init__(
+            self,
+            *,
+            _frame: types.FrameType,
+            **kwargs: ta.Any,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        self._frame = _frame
+
+    @classmethod
+    def _set_frame_builtins(
+            cls,
+            frame: types.FrameType,
+            new_builtins: dict[str, ta.Any],
+    ) -> bool:
+        return _capture._set_frame_builtins(frame, frame.f_builtins, new_builtins)  # noqa
+
+    @contextlib.contextmanager
+    def _hook_context(
+            self,
+            mod_globals: ta.MutableMapping[str, ta.Any],  # noqa
+    ) -> ta.Iterator[None]:
+        old_builtins = self._frame.f_builtins
+        old_import = old_builtins['__import__']
+        new_import = functools.partial(self._new_import, old_import)
+
+        new_builtins = dict(old_builtins)
+        new_builtins['__import__'] = new_import
+        if not self._set_frame_builtins(self._frame, new_builtins):
+            raise ImportCaptureErrors.HookError
+
+        try:
+            yield
+
+        finally:
+            if self._frame.f_builtins is not new_builtins:
+                raise ImportCaptureErrors.HookError
+
+            if not self._set_frame_builtins(self._frame, old_builtins):
+                raise ImportCaptureErrors.HookError
+
+
+#
+
+
+def _new_import_capture_hook(
+        mod_globals: ta.MutableMapping[str, ta.Any],  # noqa
+        *,
+        stack_offset: int = 0,
+        **kwargs: ta.Any,
+) -> '_ImportCaptureHook':
+    frame: types.FrameType | None = sys._getframe(1 + stack_offset)  # noqa
+    if frame is None or frame.f_globals is not mod_globals:
+        raise ImportCaptureError("Can't find importing frame")
+
+    if _capture is not None:
+        return _FrameBuiltinsImportCaptureHook(_frame=frame, **kwargs)
+
+    return _SomewhatThreadSafeGlobalBuiltinsImportCaptureHook(**kwargs)
+
+
+##
+
+
+class ImportCapture:
     class Import(ta.NamedTuple):
         spec: str
         attrs: ta.Sequence[tuple[str | None, str]]
@@ -390,11 +604,14 @@ class ImportCapture:
             self,
             mod_globals: ta.MutableMapping[str, ta.Any],
             *,
+            _hook: _ImportCaptureHook,
+
             disable: bool = False,
     ) -> None:
         super().__init__()
 
         self._mod_globals = mod_globals
+        self._hook = _hook
 
         self._disabled = disable
 
@@ -445,14 +662,12 @@ class ImportCapture:
             yield self
             return
 
-        cap = _ImportCaptureImpl()
-
-        with cap.hook_context(self._mod_globals):
+        with self._hook.hook_context(self._mod_globals):
             yield self
 
-        cap.verify_state(self._mod_globals)
+        self._hook.verify_state(self._mod_globals)
 
-        blt = cap.build_captured(
+        blt = self._hook.build_captured(
             self._mod_globals,
             collect_unreferenced=unreferenced_callback is not None or raise_unreferenced,
         )