PEP302/PEP451 import hook which rewrites asserts.
| 74 | |
| 75 | |
| 76 | class AssertionRewritingHook(importlib.abc.MetaPathFinder, importlib.abc.Loader): |
| 77 | """PEP302/PEP451 import hook which rewrites asserts.""" |
| 78 | |
| 79 | def __init__(self, config: Config) -> None: |
| 80 | self.config = config |
| 81 | try: |
| 82 | self.fnpats = config.getini("python_files") |
| 83 | except ValueError: |
| 84 | self.fnpats = ["test_*.py", "*_test.py"] |
| 85 | self.session: Session | None = None |
| 86 | self._rewritten_names: dict[str, Path] = {} |
| 87 | self._must_rewrite: set[str] = set() |
| 88 | # flag to guard against trying to rewrite a pyc file while we are already writing another pyc file, |
| 89 | # which might result in infinite recursion (#3506) |
| 90 | self._writing_pyc = False |
| 91 | self._basenames_to_check_rewrite = {"conftest"} |
| 92 | self._marked_for_rewrite_cache: dict[str, bool] = {} |
| 93 | self._session_paths_checked = False |
| 94 | |
| 95 | def set_session(self, session: Session | None) -> None: |
| 96 | self.session = session |
| 97 | self._session_paths_checked = False |
| 98 | |
| 99 | # Indirection so we can mock calls to find_spec originated from the hook during testing |
| 100 | _find_spec = importlib.machinery.PathFinder.find_spec |
| 101 | |
| 102 | def find_spec( |
| 103 | self, |
| 104 | name: str, |
| 105 | path: Sequence[str | bytes] | None = None, |
| 106 | target: types.ModuleType | None = None, |
| 107 | ) -> importlib.machinery.ModuleSpec | None: |
| 108 | if self._writing_pyc: |
| 109 | return None |
| 110 | state = self.config.stash[assertstate_key] |
| 111 | if self._early_rewrite_bailout(name, state): |
| 112 | return None |
| 113 | state.trace(f"find_module called for: {name}") |
| 114 | |
| 115 | # Type ignored because mypy is confused about the `self` binding here. |
| 116 | spec = self._find_spec(name, path) # type: ignore |
| 117 | |
| 118 | if ( |
| 119 | # the import machinery could not find a file to import |
| 120 | spec is None |
| 121 | # this is a namespace package (without `__init__.py`) |
| 122 | # there's nothing to rewrite there |
| 123 | or spec.origin is None |
| 124 | # we can only rewrite source files |
| 125 | or not isinstance(spec.loader, importlib.machinery.SourceFileLoader) |
| 126 | # if the file doesn't exist, we can't rewrite it |
| 127 | or not os.path.exists(spec.origin) |
| 128 | ): |
| 129 | return None |
| 130 | else: |
| 131 | fn = spec.origin |
| 132 | |
| 133 | if not self._should_rewrite(name, fn, state): |
no outgoing calls