Implement get_code in AssertionRewritingHook so runpy can re-run modules

AUTHORS

@@ -533,6 +533,7 @@ Yuliang Shao
 Yusuke Kadowaki
 Yutian Li
 Yuval Shimon
+yuwk
 Zac Hatfield-Dodds
 Zac Palmer Laporte
 Zach Snicker

changelog/9007.bugfix.rst

@@ -0,0 +1 @@
+The assertion-rewriting import hook now implements the ``get_code`` loader API, so that a test module can re-run itself using :func:`runpy.run_module` without crashing with ``AttributeError: 'AssertionRewritingHook' object has no attribute 'get_code'``.

src/_pytest/assertion/rewrite.py

@@ -152,6 +152,26 @@ def exec_module(self, module: types.ModuleType) -> None:
 
         self._rewritten_names[module.__name__] = fn
 
+        co = self._get_rewritten_code(fn, state)
+        exec(co, module.__dict__)
+
+    def get_code(self, name: str) -> types.CodeType | None:
+        """Return the rewritten code object for *name*, if it can be found.
+
+        This implements the optional ``get_code`` loader API, which is used by
+        :func:`runpy.run_module` among others, so that a test module can re-run
+        itself via ``runpy`` (see :issue:`9007`).
+        """
+        state = self.config.stash[assertstate_key]
+        fn = self._rewritten_names.get(name)
+        if fn is None:
+            spec = self._find_spec(name)
+            if spec is None or spec.origin is None:
+                return None
+            fn = Path(spec.origin)
+        return self._get_rewritten_code(fn, state)
+
+    def _get_rewritten_code(self, fn: Path, state: AssertionState) -> types.CodeType:
         # The requested module looks like a test file, so rewrite it. This is
         # the most magical part of the process: load the source, rewrite the
         # asserts, and load the rewritten source. We also cache the rewritten
@@ -184,7 +204,7 @@ def exec_module(self, module: types.ModuleType) -> None:
                     self._writing_pyc = False
         else:
             state.trace(f"found cached rewritten pyc for {fn}")
-        exec(co, module.__dict__)
+        return co
 
     def _early_rewrite_bailout(self, name: str, state: AssertionState) -> bool:
         """A fast way to get out of rewriting modules.

testing/test_assertrewrite.py

@@ -1318,6 +1318,73 @@ def test_meta_path():
         )
         assert pytester.runpytest().ret == 0
 
+    def test_runpy_run_module(self, pytester: Pytester) -> None:
+        """See #9007: re-running a test module with ``runpy`` should not crash."""
+        tests = pytester.mkpydir("tests")
+        tests.joinpath("test_runpy.py").write_text(
+            textwrap.dedent(
+                """
+                import runpy
+                import warnings
+
+                def test_run_module():
+                    with warnings.catch_warnings():
+                        warnings.simplefilter("ignore", RuntimeWarning)
+                        runpy.run_module("tests.test_runpy")
+                """
+            ),
+            encoding="utf-8",
+        )
+        pytester.runpytest("tests/test_runpy.py").assert_outcomes(passed=1)
+
+    def test_runpy_run_module_nested_session(self, pytester: Pytester) -> None:
+        """An outer rewrite hook remains usable while an inner session is active."""
+        pytester.makepyfile(
+            test_outer="""
+                import pytest
+                import textwrap
+
+                def test_nested_session(tmp_path):
+                    nested_test = tmp_path / "test_nested.py"
+                    nested_test.write_text(
+                        textwrap.dedent('''
+                        import runpy
+                        import sys
+
+                        from _pytest.assertion.rewrite import AssertionRewritingHook
+
+                        def test_inner():
+                            outer_loader = sys.modules["test_outer"].__spec__.loader
+                            assert isinstance(outer_loader, AssertionRewritingHook)
+                            assert outer_loader in sys.meta_path
+                            runpy.run_module("test_outer")
+                        '''),
+                        encoding="utf-8",
+                    )
+                    assert pytest.main([str(nested_test), "-q"]) == pytest.ExitCode.OK
+            """
+        )
+        pytester.runpytest().assert_outcomes(passed=1)
+
+    def test_get_code_unknown_module(self, pytestconfig, monkeypatch) -> None:
+        """The loader reports no code when the requested module cannot be found."""
+        hook = AssertionRewritingHook(pytestconfig)
+        monkeypatch.setattr(hook, "_find_spec", lambda name: None)
+
+        assert hook.get_code("unknown_module") is None
+
+    def test_get_code_unloaded_module(
+        self, pytestconfig, monkeypatch, tmp_path: Path
+    ) -> None:
+        """The loader can rewrite a resolvable module it has not executed yet."""
+        source = tmp_path / "test_unloaded.py"
+        source.write_text("assert True\n", encoding="utf-8")
+        hook = AssertionRewritingHook(pytestconfig)
+        spec = importlib.util.spec_from_file_location("test_unloaded", source)
+        monkeypatch.setattr(hook, "_find_spec", lambda name: spec)
+
+        assert hook.get_code("test_unloaded") is not None
+
     def test_write_pyc(self, pytester: Pytester, tmp_path) -> None:
         from _pytest.assertion import AssertionState
         from _pytest.assertion.rewrite import _write_pyc