diff --git a/AUTHORS b/AUTHORS index d1a2d3e7911..f861fd71d27 100644 --- a/AUTHORS +++ b/AUTHORS @@ -533,6 +533,7 @@ Yuliang Shao Yusuke Kadowaki Yutian Li Yuval Shimon +yuwk Zac Hatfield-Dodds Zac Palmer Laporte Zach Snicker diff --git a/changelog/9007.bugfix.rst b/changelog/9007.bugfix.rst new file mode 100644 index 00000000000..886232e540e --- /dev/null +++ b/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'``. diff --git a/src/_pytest/assertion/rewrite.py b/src/_pytest/assertion/rewrite.py index 362c93d7253..ff301506f5b 100644 --- a/src/_pytest/assertion/rewrite.py +++ b/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. diff --git a/testing/test_assertrewrite.py b/testing/test_assertrewrite.py index 12e12449693..85d546df1b7 100644 --- a/testing/test_assertrewrite.py +++ b/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