Add IPython support to defer_imports.console and attempt to document it.

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ repos:
       - id: check-yaml
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.3
+    rev: v0.6.4
     hooks:
       - id: ruff
         args: [--fix]

README.rst

@@ -38,19 +38,21 @@ Setup
 
 ``defer-imports`` hooks into the Python import system with a path hook. That path hook needs to be registered before code using the import-delaying context manager, ``defer_imports.until_use``, is parsed. To do that, include the following somewhere such that it will be executed before your code:
 
-.. code:: python
+.. code-block:: python
 
     import defer_imports
 
     defer_imports.install_defer_import_hook()
 
+    import your_code
+
 
 Example
 -------
 
 Assuming the path hook has been registered, you can use the ``defer_imports.until_use`` context manager to decide which imports should be deferred. For instance:
 
-.. code:: python
+.. code-block:: python
 
     import defer_imports
 
@@ -71,6 +73,77 @@ Use Cases
 -   If expensive imports are only necessary for certain code paths that won't always be taken, e.g. in subcommands in CLI tools.
 
 
+Extra: Console Usage
+--------------------
+
+``defer-imports`` works while within a regular Python REPL, as long as that work is being done in a package being imported and not with direct usage of the ``defer_imports.until_use`` context manager. To directly use the context manager, use the included interactive console:
+
+-   From the command line::
+
+        > python -m defer_imports
+        Python 3.11.9 (tags/v3.11.9:de54cf5, Apr  2 2024, 10:12:12) [MSC v.1938 64 bit (AMD64)] on win32
+        Type "help", "copyright", "credits" or "license" for more information.
+        (DeferredInteractiveConsole)
+        >>> import defer_imports
+        >>> with defer_imports.until_use:
+        ...     import typing
+        ... 
+        >>> import sys           
+        >>> "typing" in sys.modules
+        False
+        >>> typing
+        <module 'typing' from 'C:\\Users\\Tusha\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\typing.py'>
+        >>> "typing" in sys.modules
+        True
+
+-   From within a standard Python interpreter:
+
+    .. code-block:: pycon
+
+        >>> from defer_imports import console
+        >>> console.interact()
+        Python 3.11.9 (tags/v3.11.9:de54cf5, Apr  2 2024, 10:12:12) [MSC v.1938 64 bit (AMD64)] on win32
+        Type "help", "copyright", "credits" or "license" for more information.
+        (DeferredInteractiveConsole)
+        >>> import defer_imports
+        >>> with defer_imports.until_use:
+        ...     import typing
+        ... 
+        >>> import sys           
+        >>> "typing" in sys.modules
+        False
+        >>> typing
+        <module 'typing' from 'C:\\Users\\Tusha\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\typing.py'>
+        >>> "typing" in sys.modules
+        True
+
+
+Additionally, if you're using IPython in a terminal or in a Jupyter environment, there is a function you can call to ensure the context manager works there as well:
+
+.. code-block:: ipython
+
+    In [1]: from defer_imports import console
+
+    In [2]: console.instrument_ipython()
+
+    In [3]: import defer_imports
+
+    In [4]: with defer_imports.until_use:
+    ...:     import numpy
+    ...:
+
+    In [5]: import sys
+
+    In [6]: "numpy" in sys.modules
+
+    In [7]: print("numpy" in sys.modules)
+    False
+
+    In [8]: numpy
+
+    In [9]: print("numpy" in sys.modules)
+    True
+
 Features
 ========
 

src/defer_imports/_core.py

@@ -38,7 +38,12 @@ class DeferredInstrumenter(ast.NodeTransformer):
     results are assigned to custom keys in the global namespace.
     """
 
-    def __init__(self, data: SourceData, filepath: _tp.Union[_tp.StrPath, _tp.ReadableBuffer], encoding: str) -> None:
+    def __init__(
+        self,
+        data: _tp.Union[_tp.ReadableBuffer, str, ast.AST],
+        filepath: _tp.Union[_tp.StrPath, _tp.ReadableBuffer],
+        encoding: str,
+    ) -> None:
         self.data = data
         self.filepath = filepath
         self.encoding = encoding

src/defer_imports/console.py

@@ -2,17 +2,46 @@
 #
 # SPDX-License-Identifier: MIT
 
-from code import InteractiveConsole
+import __future__
 
+import ast
+import code
+import codeop
+
+from . import _typing as _tp
 from ._core import DeferredImportKey, DeferredImportProxy, DeferredInstrumenter
 
 
-__all__ = ("DeferredInteractiveConsole",)
+_features = [getattr(__future__, feat_name) for feat_name in __future__.all_feature_names]
+
+__all__ = ("DeferredInteractiveConsole", "interact", "instrument_ipython")
+
+
+class DeferredCompile(codeop.Compile):
+    """A subclass of codeop.Compile that alters the compilation process with defer_imports's AST transformer."""
+
+    def __call__(self, source: str, filename: str, symbol: str, **kwargs: object) -> _tp.CodeType:
+        flags = self.flags
+        if kwargs.get("incomplete_input", True) is False:
+            flags &= ~codeop.PyCF_DONT_IMPLY_DEDENT  # pyright: ignore
+            flags &= ~codeop.PyCF_ALLOW_INCOMPLETE_INPUT  # pyright: ignore
+        assert isinstance(flags, int)
 
+        og_ast_node = compile(source, filename, symbol, flags | ast.PyCF_ONLY_AST, True)
+        transformer = DeferredInstrumenter(source, filename, "utf-8")
+        new_ast_node = ast.fix_missing_locations(transformer.visit(og_ast_node))
 
-class DeferredInteractiveConsole(InteractiveConsole):
-    """An emulator of the interactive Python interpreter, but with defer_import's compile-time hook baked in to ensure that
-    defer_imports.until_use works as intended directly in the console.
+        codeob = compile(new_ast_node, filename, symbol, flags, True)
+        for feature in _features:
+            if codeob.co_flags & feature.compiler_flag:
+                self.flags |= feature.compiler_flag
+        return codeob
+
+
+class DeferredInteractiveConsole(code.InteractiveConsole):
+    """An emulator of the interactive Python interpreter, but with defer_import's compile-time AST transformer baked in.
+
+    This ensures that defer_imports.until_use works as intended when used directly in this console.
     """
 
     def __init__(self) -> None:
@@ -23,27 +52,38 @@ def __init__(self) -> None:
             "@DeferredImportProxy": DeferredImportProxy,
         }
         super().__init__(local_ns)
+        self.compile = codeop.CommandCompiler()
+        self.compile.compiler = DeferredCompile()
+
+
+def interact() -> None:
+    """Closely emulate the interactive Python console, but instrumented by defer_imports.
+
+    This supports direct use of the defer_imports.until_use context manager.
+    """
+
+    DeferredInteractiveConsole().interact()
+
+
+def instrument_ipython() -> None:
+    """Add defer_import's compile-time AST transformer to a currently running IPython environment.
+
+    This will ensure that defer_imports.until_use works as intended when used directly in a IPython console.
+    """
+
+    try:
+        ipython_shell: _tp.Any = get_ipython()  # pyright: ignore
+    except NameError:
+        msg = "Not currently in an IPython/Jupyter environment."
+        raise RuntimeError(msg) from None
+
+    class DeferredIPythonInstrumenter(ast.NodeTransformer):
+        def __init__(self):
+            self.actual_transformer = DeferredInstrumenter("", "<unknown>", "utf-8")
+
+        def visit(self, node: ast.AST) -> _tp.Any:
+            self.actual_transformer.data = node
+            self.actual_transformer.scope_depth = 0
+            return ast.fix_missing_locations(self.actual_transformer.visit(node))
 
-    def runsource(self, source: str, filename: str = "<input>", symbol: str = "single") -> bool:
-        try:
-            code = self.compile(source, filename, symbol)
-        except (OverflowError, SyntaxError, ValueError):
-            # Case 1: Input is incorrect.
-            self.showsyntaxerror(filename)
-            return False
-
-        if code is None:
-            # Case 2: Input is incomplete.
-            return True
-
-        # Case 3: Input is complete.
-        try:
-            tree = DeferredInstrumenter(source, filename, "utf-8").instrument(symbol)
-            code = compile(tree, filename, symbol)
-        except SyntaxError:
-            # Case 1, again.
-            self.showsyntaxerror(filename)
-            return False
-
-        self.runcode(code)
-        return False
+    ipython_shell.ast_transformers.append(DeferredIPythonInstrumenter())