Add some test cases for try-except escape hatch

bench/generate_samples.py

@@ -4,6 +4,9 @@
 from pathlib import Path
 
 
+_PYRIGHT_IGNORE_DIRECTIVES = "# pyright: reportUnusedImport=none, reportMissingTypeStubs=none"
+_GENERATED_BY_COMMENT = "# Generated by bench/generate_samples.py"
+
 # Modified from https://gist.github.com/indygreg/be1c229fa41ced5c76d912f7073f9de6.
 _STDLIB_IMPORTS = """\
 import __future__
@@ -552,8 +555,6 @@
 _INDENTED_STDLIB_IMPORTS = "".join(
     (f'{" " * 4}{line}' if line.strip() else line) for line in _STDLIB_IMPORTS.splitlines(keepends=True)
 )
-_PYRIGHT_IGNORE_DIRECTIVES = "# pyright: reportUnusedImport=none, reportMissingTypeStubs=none"
-_GENERATED_BY_COMMENT = "# Generated by bench/generate_samples.py"
 
 _CONTEXT_MANAGER_TEMPLATE = f"""\
 {_PYRIGHT_IGNORE_DIRECTIVES}

src/defer_imports/__init__.py

@@ -48,10 +48,11 @@ def _lazy_import_module(name: str, package: typing.Optional[str] = None) -> type
 
     Some types of ineligible imports:
 
-    -   from imports
-    -   submodule imports
-    -   module imports where the modules replace themselves in sys.modules during execution, e.g. collections.abc in
-        CPython 3.13
+    -   from imports (where the parent is also expected to be lazy-loaded)
+    -   submodule imports (where the parent is also expected to be lazy-loaded)
+    -   module imports where the modules replace themselves in sys.modules during execution
+        -   often done for performance reasons, like replacing onself with a c-accelerated module
+        -   e.g. collections.abc in CPython 3.13
     """
 
     # 1. Resolve the name.
@@ -98,7 +99,7 @@ def _lazy_import_module(name: str, package: typing.Optional[str] = None) -> type
 if TYPE_CHECKING:
     import ast
     import collections
-    import importlib.abc as imp_abc
+    import importlib.abc as importlib_abc
     import io
     import os
     import tokenize
@@ -107,15 +108,15 @@ def _lazy_import_module(name: str, package: typing.Optional[str] = None) -> type
     import warnings
 else:
     # fmt: off
-    ast         = _lazy_import_module("ast")
-    collections = _lazy_import_module("collections")
-    imp_abc     = _lazy_import_module("importlib.abc")
-    io          = _lazy_import_module("io")
-    os          = _lazy_import_module("os")
-    tokenize    = _lazy_import_module("tokenize")
-    types       = _lazy_import_module("types")
-    typing      = _lazy_import_module("typing")
-    warnings    = _lazy_import_module("warnings")
+    ast             = _lazy_import_module("ast")
+    collections     = _lazy_import_module("collections")
+    importlib_abc   = _lazy_import_module("importlib.abc")
+    io              = _lazy_import_module("io")
+    os              = _lazy_import_module("os")
+    tokenize        = _lazy_import_module("tokenize")
+    types           = _lazy_import_module("types")
+    typing          = _lazy_import_module("typing")
+    warnings        = _lazy_import_module("warnings")
     # fmt: on
 
 
@@ -891,7 +892,7 @@ def __init__(
         apply_all: bool,
         module_names: typing.Sequence[str],
         recursive: bool,
-        loader_class: typing.Optional[type[imp_abc.Loader]],
+        loader_class: typing.Optional[type[importlib_abc.Loader]],
     ) -> None:
         self.apply_all = apply_all
         self.module_names = module_names
@@ -955,7 +956,7 @@ def install_import_hook(
     apply_all: bool = False,
     module_names: typing.Sequence[str] = (),
     recursive: bool = False,
-    loader_class: typing.Optional[type[imp_abc.Loader]] = None,
+    loader_class: typing.Optional[type[importlib_abc.Loader]] = None,
 ) -> ImportHookContext:
     r"""Install defer_imports's import hook if it isn't already installed, and optionally configure it. Must be called
     before using defer_imports.until_use.
@@ -1128,8 +1129,7 @@ def _resolve(self) -> None:
         module_vars = vars(module)
         for attr_key, attr_val in vars(proxy).items():
             if isinstance(attr_val, _DeferredImportProxy) and not hasattr(module, attr_key):
-                # This could have used setattr() if pypy didn't normalize the attr key type to str, so we resort to
-                # direct placement in the module's __dict__ to avoid that.
+                # NOTE: This doesn't use setattr() because pypy normalizes the attr key type to str.
                 module_vars[_DeferredImportKey(attr_key, attr_val)] = attr_val
 
                 # Change the namespaces as well to make sure nested proxies are replaced in the right place.

tests/test_defer_imports.py

@@ -68,10 +68,10 @@ def temp_cache_module(name: str, module: types.ModuleType):
 def better_key_repr(monkeypatch: pytest.MonkeyPatch):
     """Replace defer_imports._comptime._DeferredImportKey.__repr__ with a more verbose version for all tests."""
 
-    def _verbose_repr(self) -> str:  # pyright: ignore  # noqa: ANN001
-        return f"<key for {str.__repr__(self)} import>"  # pyright: ignore
+    def verbose_repr(self: object) -> str:
+        return f"<key for {super(type(self), self).__repr__()} import>"
 
-    monkeypatch.setattr("defer_imports._DeferredImportKey.__repr__", _verbose_repr)  # pyright: ignore [reportUnknownArgumentType]
+    monkeypatch.setattr("defer_imports._DeferredImportKey.__repr__", verbose_repr)
 
 
 # endregion
@@ -524,6 +524,82 @@ def do_the_thing(a: int) -> int:
 """,
             id="still instruments imports in defer_imports.until_use with block",
         ),
+        pytest.param(
+            """\
+try:
+    import foo
+except:
+    pass
+""",
+            """\
+import defer_imports
+from defer_imports import _DeferredImportKey as @_DeferredImportKey, _DeferredImportProxy as @_DeferredImportProxy
+try:
+    import foo
+except:
+    pass
+del @_DeferredImportKey, @_DeferredImportProxy
+""",
+            id="escape hatch: try",
+        ),
+        pytest.param(
+            """\
+try:
+    raise Exception
+except:
+    import foo
+""",
+            """\
+import defer_imports
+from defer_imports import _DeferredImportKey as @_DeferredImportKey, _DeferredImportProxy as @_DeferredImportProxy
+try:
+    raise Exception
+except:
+    import foo
+del @_DeferredImportKey, @_DeferredImportProxy
+""",
+            id="escape hatch: except",
+        ),
+        pytest.param(
+            """\
+try:
+    print('hi')
+except:
+    print('error')
+else:
+    import foo
+""",
+            """\
+import defer_imports
+from defer_imports import _DeferredImportKey as @_DeferredImportKey, _DeferredImportProxy as @_DeferredImportProxy
+try:
+    print('hi')
+except:
+    print('error')
+else:
+    import foo
+del @_DeferredImportKey, @_DeferredImportProxy
+""",
+            id="escape hatch: else",
+        ),
+        pytest.param(
+            """\
+try:
+    pass
+finally:
+    import foo
+""",
+            """\
+import defer_imports
+from defer_imports import _DeferredImportKey as @_DeferredImportKey, _DeferredImportProxy as @_DeferredImportProxy
+try:
+    pass
+finally:
+    import foo
+del @_DeferredImportKey, @_DeferredImportProxy
+""",
+            id="escape hatch: finally",
+        ),
     ],
 )
 def test_module_instrumentation(before: str, after: str):
@@ -576,7 +652,6 @@ def test_regular_import(tmp_path: Path):
     spec, module, _ = create_sample_module(tmp_path, source)
     assert spec.loader
     spec.loader.exec_module(module)
-    print(repr(vars(module)))
 
     expected_inspect_repr = "<key for 'inspect' import>: <proxy for 'import inspect'>"
     assert expected_inspect_repr in repr(vars(module))
@@ -904,8 +979,6 @@ def test_top_level_and_submodules_1(tmp_path: Path):
 
 def test_top_level_and_submodules_2(tmp_path: Path):
     source = """\
-from pprint import pprint
-
 import defer_imports
 
 with defer_imports.until_use:
@@ -1116,17 +1189,23 @@ def Y2():
 def test_import_stdlib():
     """Test that defer_imports.until_use works when wrapping imports for most of the stdlib."""
 
-    # The finder for tests.sample_stdlib_imports is already cached, so we need to temporarily reset that cache.
-    _temp_cache = dict(sys.path_importer_cache)
-    sys.path_importer_cache.clear()
+    _missing: Any = object()
 
-    with install_import_hook(uninstall_after=True):
-        import tests.sample_stdlib_imports
+    # The path finder for the tests directory is already cached, so we need to temporarily reset that entry.
+    _tests_path = str(Path(__file__).parent)
+    _temp_cache = sys.path_importer_cache.pop(_tests_path, _missing)
 
-    assert tests.sample_stdlib_imports
+    try:
+        with install_import_hook(uninstall_after=True):
+            import tests.sample_stdlib_imports
 
-    # Revert changes to the path finder cache.
-    sys.path_importer_cache = _temp_cache
+        # Sample-check the __future__ import.
+        expected_future_import = "<key for '__future__' import>: <proxy for 'import __future__'>"
+        assert expected_future_import in repr(vars(tests.sample_stdlib_imports))
+    finally:
+        # Revert changes to the path finder cache.
+        if _temp_cache is not _missing:
+            sys.path_importer_cache[_tests_path] = _temp_cache
 
 
 @pytest.mark.skip(reason="Leaking patch problem is currently out of scope.")