Implement PEP 302 optional get_code loader method

src/_pytest/assertion/rewrite.py

@@ -80,6 +80,7 @@ def __init__(self, config: Config) -> None:
         self._basenames_to_check_rewrite = {"conftest"}
         self._marked_for_rewrite_cache: dict[str, bool] = {}
         self._session_paths_checked = False
+        self.fn: str | None = None
 
     def set_session(self, session: Session | None) -> None:
         self.session = session
@@ -126,7 +127,7 @@ def find_spec(
         ):
             return None
         else:
-            fn = spec.origin
+            self.fn = fn = spec.origin
 
         if not self._should_rewrite(name, fn, state):
             return None
@@ -143,14 +144,11 @@ def create_module(
     ) -> types.ModuleType | None:
         return None  # default behaviour is fine
 
-    def exec_module(self, module: types.ModuleType) -> None:
-        assert module.__spec__ is not None
-        assert module.__spec__.origin is not None
-        fn = Path(module.__spec__.origin)
+    def get_code(self, fullname: str) -> types.CodeType:
+        assert self.fn is not None
+        fn = Path(self.fn)
         state = self.config.stash[assertstate_key]
 
-        self._rewritten_names[module.__name__] = fn
-
         # 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
@@ -183,7 +181,15 @@ 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 exec_module(self, module: types.ModuleType) -> None:
+        module_name = module.__name__
+
+        self._rewritten_names[module_name] = fn
+
+        exec(self.get_code(module_name), module.__dict__)
 
     def _early_rewrite_bailout(self, name: str, state: AssertionState) -> bool:
         """A fast way to get out of rewriting modules.
@@ -1213,4 +1219,4 @@ def get_cache_dir(file_path: Path) -> Path:
         return Path(sys.pycache_prefix) / Path(*file_path.parts[1:-1])
     else:
         # classic pycache directory
-        return file_path.parent / "__pycache__"
+        return file_path.parent / "__pycache__"