Lift out _SubmoduleAutoImporterModule and register it to guarded_eval.

anntzer · anntzer · commit 2bf460305e8d · 2023-04-22T12:20:55.000+02:00
diff --git a/README.rst b/README.rst
@@ -64,9 +64,12 @@ As usual, install using pip:
 
 Then, append the output of ``python -m ipython_autoimport`` to the
 ``ipython_config.py`` file in the directory printed by ``ipython profile
-locate`` (typically ``~/.ipython/profile_default/``).
+locate`` (typically ``~/.ipython/profile_default/``).  If you don't have such a
+file at all, first create it with ``ipython profile create``.
 
-If you don't have such a file at all, you can use ``ipython profile create``.
+Note that upon loading, ``ipython_autoimport`` will register its submodule
+auto-importer to IPython's "limited evalutation" completer policy (on IPython
+versions that support it).
 
 Run tests with ``pytest``.
 
@@ -87,6 +90,6 @@ the builtins dict instead, but that seems a bit too invasive...).
 
 When using Jedi autocompletion (the default if Jedi is installed as of IPython
 7.2), trying to tab-complete not-yet-imported global names to trigger an import
-fails, because Jedi purposefully converts the global dict to a namespace
+failure, because Jedi purposefully converts the global dict to a namespace
 object and looks up attributes using ``getattr_static``.  Jedi can be disabled
 by adding ``c.Completer.use_jedi = False`` to the ``ipython_config.py`` file.
diff --git a/lib/ipython_autoimport.py b/lib/ipython_autoimport.py
@@ -15,9 +15,9 @@
 import token
 from types import ModuleType
 
+import IPython.core
 from IPython.core import magic
 from IPython.core.error import UsageError
-from IPython.core.magic import register_line_magic
 from IPython.core.magic_arguments import (
     argument, magic_arguments, parse_argstring)
 from IPython.core.magics.execution import ExecutionMagics
@@ -93,53 +93,54 @@ def _report(ipython, msg):
     print("{}Autoimport:{} {}".format(cs[token.NUMBER], cs["normal"], msg))
 
 
-def _make_submodule_autoimporter_module(ipython, module):
-    """
-    Return a module sub-instance that automatically imports submodules.
+class _SubmoduleAutoImporterModule(ModuleType):
+    # __module and __ipython are set externally to not modify the constructor.
 
-    Implemented as a factory function to close over the real module.
-    """
+    @property
+    def __dict__(self):
+        return self.__module.__dict__
 
-    if not hasattr(module, "__path__"):  # We only need to wrap packages.
-        return module
+    # Overriding __setattr__ is needed even when __dict__ is overridden.
+    def __setattr__(self, name, value):
+        setattr(self.__module, name, value)
 
-    class SubmoduleAutoImporterModule(ModuleType):
-        @property
-        def __dict__(self):
-            return module.__dict__
+    def __getattr__(self, name):
+        try:
+            value = getattr(self.__module, name)
+            if isinstance(value, ModuleType):
+                value = _make_submodule_autoimporter_module(
+                    self.__ipython, value)
+            return value
+        except AttributeError:
+            import_target = "{}.{}".format(self.__name__, name)
+            try:
+                submodule = importlib.import_module(import_target)
+            except getattr(builtins, "ModuleNotFoundError", ImportError):
+                pass  # Py<3.6.
+            else:
+                _report(self.__ipython, "import {}".format(import_target))
+                return _make_submodule_autoimporter_module(
+                    self.__ipython, submodule)
+            raise  # Raise AttributeError without chaining ImportError.
 
-        # Overriding __setattr__ is needed even when __dict__ is overridden.
-        def __setattr__(self, name, value):
-            setattr(module, name, value)
 
-        def __getattr__(self, name):
-            try:
-                value = getattr(module, name)
-                if isinstance(value, ModuleType):
-                    value = _make_submodule_autoimporter_module(ipython,
-                                                                value)
-                return value
-            except AttributeError:
-                import_target = "{}.{}".format(self.__name__, name)
-                try:
-                    submodule = importlib.import_module(import_target)
-                except getattr(builtins, "ModuleNotFoundError",
-                               ImportError):  # Py<3.6.
-                    pass
-                else:
-                    _report(ipython, "import {}".format(import_target))
-                    return _make_submodule_autoimporter_module(
-                        ipython, submodule)
-                raise  # Raise AttributeError without chaining ImportError.
-
-    sai_module = SubmoduleAutoImporterModule(module.__name__)
-    # Apparently, `module?` does not trigger descriptors, so we need to
-    # set the docstring explicitly (on the instance, not on the class).
-    # Then then only difference in the output of `module?` becomes the type
-    # (`SubmoduleAutoImportModule` instead of `module`), which we should keep
-    # for clarity.
-    ModuleType.__setattr__(sai_module, "__doc__", module.__doc__)
-    return sai_module
+def _make_submodule_autoimporter_module(ipython, module):
+    """Return a module sub-instance that automatically imports submodules."""
+    if not hasattr(module, "__path__"):  # We only need to wrap packages.
+        return module
+    saim = _SubmoduleAutoImporterModule(module.__name__)
+    for k, v in [
+            ("_SubmoduleAutoImporterModule__module", module),
+            ("_SubmoduleAutoImporterModule__ipython", ipython),
+            # Apparently, `module?` does not trigger descriptors, so we need to
+            # set the docstring explicitly (on the instance, not on the class).
+            # Then then only difference in the output of `module?` becomes the
+            # type (`SubmoduleAutoImportModule` instead of `module`), which we
+            # should keep for clarity.
+            ("__doc__", module.__doc__),
+    ]:
+        ModuleType.__setattr__(saim, k, v)
+    return saim
 
 
 class _AutoImporterMap(dict):
@@ -231,6 +232,9 @@ def _install_namespace(ipython):
     # (both with and without Jedi).
     ipython.user_ns = ipython.Completer.namespace = (
         _AutoImporterMap(ipython))
+    if hasattr(IPython.core, "guarded_eval"):
+        (IPython.core.guarded_eval.EVALUATION_POLICIES["limited"]
+         .allowed_getattr.add(_SubmoduleAutoImporterModule))
 
 
 def _uninstall_namespace(ipython):
@@ -266,7 +270,7 @@ def autoimport(arg):
 def load_ipython_extension(ipython):
     _install_namespace(ipython)
     ipython.register_magics(_PatchedMagics)  # Add warning to timing magics.
-    register_line_magic(autoimport)
+    magic.register_line_magic(autoimport)
 
 
 def unload_ipython_extension(ipython):
