diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d41f984..1b8eb81 100644 --- a/.pre-commit-config.yaml +++ b/.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] diff --git a/README.rst b/README.rst index 49a6670..d686962 100644 --- a/README.rst +++ b/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 + + >>> "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 + + >>> "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 ======== diff --git a/src/defer_imports/_core.py b/src/defer_imports/_core.py index 3c8d5cb..93732c5 100644 --- a/src/defer_imports/_core.py +++ b/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 diff --git a/src/defer_imports/console.py b/src/defer_imports/console.py index 676f30a..d575fc8 100644 --- a/src/defer_imports/console.py +++ b/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("", "", "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 = "", 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())