1"""
2Loading a test module.
3
4A test module is loaded by Python's import machinery, through a loader that
5differs from the ordinary one in one step. `source_to_code` is where a file's
6source becomes a code object, and here it checks the file's imports and
7rewrites its asserts on the way. Everything else is the standard library's:
8the module has a spec, a loader, a file and a cache, so `inspect`, `pickle`,
9tracebacks and debuggers treat it as they treat any module.
10
11The runner builds the spec for each file it has already decided is a test
12file. Nothing is installed on `sys.meta_path`: there is no import to
13intercept, because nothing imports a test module by name.
14
15The rewritten bytecode is cached beside the ordinary cache, under a name an
16ordinary import never looks for. See `cache_path_for`.
17"""
18
19import ast
20import dataclasses
21import functools
22import hashlib
23import importlib.machinery
24import importlib.util
25import marshal
26import os
27import sys
28import time
29import types
30from dataclasses import dataclass
31from pathlib import Path
32
33from ..definition import TestDefinitionError
34from . import assertions, problems
35from .layout import TEST_MODULES_PACKAGE, Layout
36from .problems import CantBeRunAsWritten, ProblemsInAFile, tests_as_written
37
38__all__ = []
39
40_CACHE_TAG_PREFIX = "plaintest"
41
42
43@dataclass(frozen=True, kw_only=True)
44class RewriterWork:
45 """
46 The test files the loader has handled in this process: how many were
47 rewritten and compiled, how many were read back from the cache, and how
48 long each took. A cold cache is the difference between the two.
49 """
50
51 rewritten_files: int = 0
52 rewritten_seconds: float = 0.0
53 cached_files: int = 0
54 cached_seconds: float = 0.0
55
56 def since(self, before: RewriterWork) -> RewriterWork:
57 return RewriterWork(
58 rewritten_files=self.rewritten_files - before.rewritten_files,
59 rewritten_seconds=self.rewritten_seconds - before.rewritten_seconds,
60 cached_files=self.cached_files - before.cached_files,
61 cached_seconds=self.cached_seconds - before.cached_seconds,
62 )
63
64
65_work = RewriterWork()
66
67
68def rewriter_work() -> RewriterWork:
69 """The work done so far. Take one before and one after to count a run's."""
70 return _work
71
72
73# A bytecode file's header: the interpreter's magic number, a flags field
74# (zero for a file checked against its source's time and size), then the
75# time and the size, four bytes each.
76_NO_FLAGS = (0).to_bytes(4, "little")
77_HEADER_LENGTH = 16
78
79
80def load_test_module(path: Path, *, layout: Layout) -> types.ModuleType:
81 """The module for one test file, loaded if it hasn't been."""
82 name = _module_name_for(path, layout=layout)
83 loaded = sys.modules.get(name)
84 # The name says where a file is under the root, so two roots have files
85 # of the same name. Only the same file is the same module.
86 if loaded is not None and getattr(loaded, "__file__", None) == str(path):
87 return loaded
88
89 loader = TestModuleLoader(name, str(path), layout=layout)
90 spec = importlib.util.spec_from_file_location(name, path, loader=loader)
91 assert spec is not None
92 # Where this module's bytecode really is, not where an ordinary import
93 # of the same file would have put it.
94 spec.cached = loader.cache_path
95 module = importlib.util.module_from_spec(spec)
96
97 _add_to_its_packages(module)
98 sys.modules[name] = module
99 try:
100 loader.exec_module(module)
101 except ModuleNotFoundError as error:
102 sys.modules.pop(name, None)
103 problems = imports_of_a_helper_by_part_of_its_path(
104 path, missing=error.name, layout=layout
105 )
106 if problems:
107 raise ImportsAnotherWay(problems, layout=layout) from error
108 raise
109 except BaseException:
110 sys.modules.pop(name, None)
111 raise
112 return module
113
114
115def _module_name_for(path: Path, *, layout: Layout) -> str:
116 if path.is_relative_to(layout.root):
117 relative = path.relative_to(layout.root)
118 return ".".join([TEST_MODULES_PACKAGE, *relative.with_suffix("").parts])
119 return f"{TEST_MODULES_PACKAGE}.{path.stem}"
120
121
122def _add_to_its_packages(module: types.ModuleType) -> None:
123 """
124 Make the packages a test module's name says it is in, and put each one
125 in the one above it.
126
127 `plain_tests.public.test_views` names two packages that exist nowhere
128 on disk. Anything that finds a module by its name imports the packages
129 first (unpickling an instance of a class the test file defines is one),
130 so they have to be there. They are empty: a package with no search
131 path, which nothing can be imported from.
132 """
133 child = module
134 name = module.__name__
135 while "." in name:
136 name, _, child_name = name.rpartition(".")
137 package = sys.modules.get(name)
138 if package is None:
139 package = importlib.util.module_from_spec(
140 importlib.machinery.ModuleSpec(name, loader=None, is_package=True)
141 )
142 sys.modules[name] = package
143 setattr(package, child_name, child)
144 child = package
145
146
147class TestModuleLoader(importlib.machinery.SourceFileLoader):
148 def __init__(self, fullname: str, path: str, *, layout: Layout) -> None:
149 super().__init__(fullname, path)
150 self.layout = layout
151 self.cache_path = cache_path_for(
152 path, refused_import_name=layout.refused_import_name
153 )
154
155 def source_to_code(self, data: bytes, path: str) -> types.CodeType: # ty: ignore[invalid-method-override]
156 source = importlib.util.decode_source(data)
157 tree = ast.parse(source, filename=path)
158
159 # A file that can't be run says everything reading it shows, about
160 # its imports and about its tests. Run, it would stop at the first
161 # of them, and each would take a run of its own to find.
162 imports = import_problems(tree, layout=self.layout, path=Path(path))
163 tests = tests_as_written(tree)
164
165 if imports:
166 raise ImportsAnotherWay(imports, tests=tests, layout=self.layout)
167
168 if tests.stops_at_a_decorator:
169 raise CantBeRunAsWritten(tests)
170
171 # Anything else wrong with its tests is found by running it and
172 # looking at them, which sees more than reading does.
173
174 tree = assertions.rewrite_asserts(tree, source=source)
175 # dont_inherit: a test module is compiled with its own __future__
176 # statements and nothing else. Without it, compile() would also apply
177 # whatever compiler flags are in effect in this file — none today,
178 # but a test module's semantics shouldn't depend on that staying true.
179 return compile(tree, path, "exec", dont_inherit=True)
180
181 def get_code(self, fullname: str) -> types.CodeType:
182 """
183 The module's code: from the cache when the cache was written from
184 the source as it is now, and from the source otherwise.
185
186 This is what the standard loader's `get_code` does, with the cache
187 in a different place. It can't be told where to look, so the steps
188 are written out here.
189 """
190 global _work
191
192 started = time.perf_counter()
193 source_path = self.get_filename(fullname)
194 stats = self.path_stats(source_path)
195 # A bytecode file has four bytes for each, as the standard one does.
196 source_time = int(stats["mtime"]) & 0xFFFFFFFF
197 source_size = stats["size"] & 0xFFFFFFFF
198 header = (
199 importlib.util.MAGIC_NUMBER
200 + _NO_FLAGS
201 + source_time.to_bytes(4, "little")
202 + source_size.to_bytes(4, "little")
203 )
204
205 if self.cache_path is not None:
206 try:
207 cached = self.get_data(self.cache_path)
208 except OSError:
209 pass # nothing cached yet
210 else:
211 if cached[:_HEADER_LENGTH] == header:
212 try:
213 code = marshal.loads(cached[_HEADER_LENGTH:])
214 except EOFError, ValueError, TypeError:
215 pass # cut short or damaged: compile it again
216 else:
217 _work = dataclasses.replace(
218 _work,
219 cached_files=_work.cached_files + 1,
220 cached_seconds=_work.cached_seconds
221 + (time.perf_counter() - started),
222 )
223 return code
224
225 code = self.source_to_code(self.get_data(source_path), source_path)
226
227 if self.cache_path is not None and not sys.dont_write_bytecode:
228 # Makes the directory, writes the file in one step, and lets a
229 # directory that can't be written to pass without a word, as it
230 # does for any module's bytecode.
231 self.set_data(self.cache_path, header + marshal.dumps(code))
232 _remove_caches_left_by_other_rewriters(Path(self.cache_path))
233 _work = dataclasses.replace(
234 _work,
235 rewritten_files=_work.rewritten_files + 1,
236 rewritten_seconds=_work.rewritten_seconds + (time.perf_counter() - started),
237 )
238 return code
239
240
241def cache_path_for(source_path: str, *, refused_import_name: str | None) -> str | None:
242 """
243 Where a test file's rewritten bytecode is kept, or None if this
244 interpreter keeps no bytecode.
245
246 It goes where Python would put the file's bytecode (`__pycache__`, or
247 under `PYTHONPYCACHEPREFIX`), with a tag of its own in the name, in the
248 place Python puts `opt-1` and `opt-2`:
249
250 __pycache__/test_views.cpython-314.opt-plaintest3f2a9c1b7d04.pyc
251
252 An ordinary import of the file reads `test_views.cpython-314.pyc`, so it
253 never gets bytecode with rewritten asserts in it.
254
255 The tag changes when the rewriter does. What a cached file holds depends
256 on the source, which its header is checked against, and on the code that
257 rewrote it, which the tag stands for.
258 """
259 tag = _CACHE_TAG_PREFIX + _what_rewrites(refused_import_name=refused_import_name)
260 try:
261 return importlib.util.cache_from_source(source_path, optimization=tag)
262 except NotImplementedError:
263 return None
264
265
266@functools.cache
267def _what_rewrites(*, refused_import_name: str | None) -> str:
268 """
269 Twelve characters that change whenever a test module would be compiled
270 differently from the same source: when the rewriter changes, when this
271 module changes, and when a different import would be refused.
272
273 The last matters because a file is only cached once its imports have
274 passed the check, and what the check refuses depends on the layout.
275 """
276 digest = hashlib.sha256()
277 for module in (assertions, problems, sys.modules[__name__]):
278 assert module.__file__ is not None
279 digest.update(Path(module.__file__).read_bytes())
280 digest.update(repr(refused_import_name).encode())
281 return digest.hexdigest()[:12]
282
283
284def _remove_caches_left_by_other_rewriters(cache_path: Path) -> None:
285 """
286 Remove what earlier versions of the rewriter cached for the same file.
287 Each one has a tag nothing will ask for again.
288 """
289 # test_views.cpython-314.opt-plaintest3f2a9c1b7d04.pyc
290 # -> test_views.cpython-314.opt-plaintest
291 until_the_tag, _, _ = cache_path.name.rpartition(_CACHE_TAG_PREFIX)
292 for other in cache_path.parent.glob(f"{until_the_tag}{_CACHE_TAG_PREFIX}*.pyc"):
293 if other != cache_path:
294 other.unlink(missing_ok=True)
295
296
297def import_problems(tree: ast.Module, *, layout: Layout, path: Path) -> list[str]:
298 """
299 Imports in the test module at `path` that reach a helper module some way
300 other than by its path from the tests directory.
301
302 A relative import is refused: the packages a test module is in are
303 empty, so there is nothing beside it to import. An import through the
304 tests directory's own name (`tests.helpers`) works only when the command
305 runs from the directory above it, and loads a second copy of a module
306 that something else imported as `helpers`.
307 """
308 directories = _directories_a_file_is_under(path, layout=layout)
309
310 problems = []
311 for node in ast.walk(tree):
312 if isinstance(node, ast.ImportFrom):
313 module = node.module or ""
314 names = ", ".join(
315 f"{alias.name} as {alias.asname}" if alias.asname else alias.name
316 for alias in node.names
317 )
318 if node.level > 0:
319 written = f"from {'.' * node.level}{module} import {names}"
320 # One dot is the file's own directory, and each dot more is
321 # a directory further up.
322 up = node.level - 1
323 beside = directories[: len(directories) - up] if up else directories
324 if up > len(directories):
325 beside = ()
326 from_module = ".".join([*beside, *([module] if module else [])])
327 elif module.split(".")[0] == layout.refused_import_name:
328 written = f"from {module} import {names}"
329 from_module = module.partition(".")[2]
330 else:
331 continue
332 if from_module:
333 corrected = f"from {from_module} import {names}"
334 else:
335 corrected = f"import {names}"
336 elif isinstance(node, ast.Import):
337 through = [
338 alias.name
339 for alias in node.names
340 if "." in alias.name
341 and alias.name.split(".")[0] == layout.refused_import_name
342 ]
343 if not through:
344 continue
345 written = f"import {through[0]}"
346 corrected = f"import {through[0].partition('.')[2]}"
347 else:
348 continue
349 problems.append(f"line {node.lineno}: `{written}` should be `{corrected}`")
350 return problems
351
352
353def _directories_a_file_is_under(path: Path, *, layout: Layout) -> tuple[str, ...]:
354 """
355 The directories between the tests directory and a file in it:
356 `("billing", "refunds")` for `tests/billing/refunds/test_partial.py`,
357 and none for a file in the tests directory itself.
358 """
359 directory = path.parent
360 if directory.is_relative_to(layout.helper_directory):
361 return directory.relative_to(layout.helper_directory).parts
362 return ()
363
364
365def imports_of_a_helper_by_part_of_its_path(
366 path: Path, *, missing: str | None, layout: Layout
367) -> list[str]:
368 """
369 What to write, when the test file at `path` couldn't import `missing`
370 and a module by that name is further down in the tests directory.
371
372 `tests/billing/helpers.py` is `billing.helpers`. A test file beside it
373 that writes `from helpers import charge` finds nothing, or finds
374 `tests/helpers.py`, which is another module. Only the first can be told
375 from here, and it is told when it happens: what is in the tests
376 directory can change without the test file changing.
377 """
378 if not missing:
379 return []
380 top_name = missing.split(".")[0]
381
382 where_it_is = []
383 for directory, directory_names, file_names in os.walk(layout.helper_directory):
384 directory_names[:] = sorted(
385 name
386 for name in directory_names
387 if not name.startswith(".") and name != "__pycache__"
388 )
389 if Path(directory) == layout.helper_directory:
390 continue
391 if f"{top_name}.py" in file_names or top_name in directory_names:
392 under = Path(directory).relative_to(layout.helper_directory).parts
393 where_it_is.append(".".join(under))
394 if not where_it_is:
395 return []
396
397 try:
398 tree = ast.parse(path.read_text(), filename=str(path))
399 except OSError, SyntaxError, ValueError:
400 return []
401
402 problems = []
403 for node in ast.walk(tree):
404 if isinstance(node, ast.ImportFrom) and node.level == 0:
405 module = node.module or ""
406 if module.split(".")[0] != top_name:
407 continue
408 names = ", ".join(
409 f"{alias.name} as {alias.asname}" if alias.asname else alias.name
410 for alias in node.names
411 )
412 written = f"from {module} import {names}"
413 corrected = [
414 f"from {under}.{module} import {names}" for under in where_it_is
415 ]
416 elif isinstance(node, ast.Import):
417 by_that_name = [
418 alias for alias in node.names if alias.name.split(".")[0] == top_name
419 ]
420 if not by_that_name:
421 continue
422 alias = by_that_name[0]
423 written = f"import {alias.name}"
424 if alias.asname:
425 written += f" as {alias.asname}"
426 corrected = [
427 f"import {under}.{alias.name} as {alias.asname or top_name}"
428 for under in where_it_is
429 ]
430 else:
431 continue
432 either = " or ".join(f"`{one}`" for one in corrected)
433 problems.append(f"line {node.lineno}: `{written}` should be {either}")
434 return problems
435
436
437class ImportsAnotherWay(TestDefinitionError):
438 """
439 A test file imports a helper module some way other than by its path
440 from the tests directory. Each problem says what to write, so a run that
441 finds them in many files says why once.
442 """
443
444 def __init__(
445 self,
446 problems: list[str],
447 *,
448 layout: Layout,
449 tests: ProblemsInAFile | None = None,
450 ) -> None:
451 self.problems = problems
452 # What reading the file showed about its tests.
453 self.tests = tests or ProblemsInAFile()
454 at_the_top = layout.shown(layout.helper_directory / "helpers.py")
455 further_down = layout.shown(layout.helper_directory / "billing" / "helpers.py")
456 self.why = (
457 "A helper module is imported by its path from the tests directory,\n"
458 "whichever directory the test file is in and wherever the command\n"
459 f"runs from: {at_the_top} is `helpers`, and\n"
460 f"{further_down} is `billing.helpers`."
461 )
462 super().__init__(f"{self.what_is_wrong()}\n\n{self.why}")
463
464 def what_is_wrong(self) -> str:
465 listed = "\n".join(f" {problem}" for problem in self.problems)
466 return f"These imports can't be used in a test file:\n\n{listed}"