|
| 1 | +""" |
| 2 | +Static extraction of the test code behind an image-comparison test |
| 3 | +
|
| 4 | +An image-comparison test identifies its image with a string literal, |
| 5 | +``assert p == "params"`` (see ``tests/conftest.py``). Given that name, |
| 6 | +this module recovers a readable snippet: the test function together |
| 7 | +with the class-level and module-level assignments it relies on. |
| 8 | +""" |
| 9 | + |
| 10 | +from __future__ import annotations |
| 11 | + |
| 12 | +import ast |
| 13 | +from dataclasses import dataclass |
| 14 | +from functools import lru_cache |
| 15 | +from typing import TYPE_CHECKING |
| 16 | + |
| 17 | +if TYPE_CHECKING: |
| 18 | + from pathlib import Path |
| 19 | + |
| 20 | +__all__ = ["CodeSnippet", "extract_test_code", "get_test_code"] |
| 21 | + |
| 22 | +FunctionNode = ast.FunctionDef | ast.AsyncFunctionDef |
| 23 | + |
| 24 | + |
| 25 | +@dataclass(frozen=True) |
| 26 | +class CodeSnippet: |
| 27 | + """ |
| 28 | + Test code that produced an image, ready for display |
| 29 | +
|
| 30 | + Not named ``Test*`` so pytest never tries to collect it. |
| 31 | + """ |
| 32 | + |
| 33 | + source: str |
| 34 | + """Assembled snippet: setup assignments, then the test function""" |
| 35 | + |
| 36 | + lineno: int |
| 37 | + """Line in the test file where the test function starts""" |
| 38 | + |
| 39 | + |
| 40 | +def extract_test_code(source: str, name: str) -> CodeSnippet | None: |
| 41 | + """ |
| 42 | + Extract the code of the test that creates image `name` |
| 43 | +
|
| 44 | + Parameters |
| 45 | + ---------- |
| 46 | + source : |
| 47 | + Contents of a test file. |
| 48 | + name : |
| 49 | + Identifier of the test image, i.e. the string the plot is |
| 50 | + compared against. |
| 51 | +
|
| 52 | + Returns |
| 53 | + ------- |
| 54 | + : |
| 55 | + The snippet, or None when `source` does not parse or no test |
| 56 | + function mentions `name`. |
| 57 | + """ |
| 58 | + try: |
| 59 | + tree = ast.parse(source) |
| 60 | + except SyntaxError: |
| 61 | + return None |
| 62 | + return _extract(source, tree, name) |
| 63 | + |
| 64 | + |
| 65 | +def get_test_code(test_file: Path, name: str) -> CodeSnippet | None: |
| 66 | + """ |
| 67 | + Extract the code of the test in `test_file` that creates image `name` |
| 68 | +
|
| 69 | + Returns None when the file is missing or unparseable. Files are |
| 70 | + read and parsed once per process. |
| 71 | + """ |
| 72 | + parsed = _parse_file(test_file) |
| 73 | + if parsed is None: |
| 74 | + return None |
| 75 | + source, tree = parsed |
| 76 | + return _extract(source, tree, name) |
| 77 | + |
| 78 | + |
| 79 | +@lru_cache(maxsize=None) |
| 80 | +def _parse_file(test_file: Path) -> tuple[str, ast.Module] | None: |
| 81 | + try: |
| 82 | + source = test_file.read_text(encoding="utf-8") |
| 83 | + return source, ast.parse(source) |
| 84 | + except (OSError, SyntaxError): |
| 85 | + return None |
| 86 | + |
| 87 | + |
| 88 | +def _extract(source: str, tree: ast.Module, name: str) -> CodeSnippet | None: |
| 89 | + lines = source.split("\n") |
| 90 | + found = _find_test_function(tree, name) |
| 91 | + if found is None: |
| 92 | + return None |
| 93 | + func, cls = found |
| 94 | + |
| 95 | + loaded = _loaded_names(func) |
| 96 | + chunks = [ |
| 97 | + _segment(lines, stmt) |
| 98 | + for targets, stmt in _module_assignments(tree) |
| 99 | + if targets & loaded |
| 100 | + ] |
| 101 | + |
| 102 | + if cls is None: |
| 103 | + chunks.append(_segment(lines, func)) |
| 104 | + else: |
| 105 | + class_chunks = [ |
| 106 | + _segment(lines, stmt) |
| 107 | + for stmt in cls.body |
| 108 | + if isinstance(stmt, (ast.Assign, ast.AnnAssign)) |
| 109 | + ] |
| 110 | + body = "\n\n".join([*class_chunks, _segment(lines, func)]) |
| 111 | + chunks.append(f"class {cls.name}:\n{body}") |
| 112 | + |
| 113 | + return CodeSnippet(source="\n\n".join(chunks), lineno=func.lineno) |
| 114 | + |
| 115 | + |
| 116 | +def _find_test_function( |
| 117 | + tree: ast.Module, name: str |
| 118 | +) -> tuple[FunctionNode, ast.ClassDef | None] | None: |
| 119 | + """ |
| 120 | + The test function that identifies its image as `name` |
| 121 | +
|
| 122 | + A function comparing a plot to `name` wins; a function that merely |
| 123 | + contains the string (e.g. in a list of names) is the fallback. |
| 124 | + """ |
| 125 | + constant_match = None |
| 126 | + for func, cls in _iter_functions(tree): |
| 127 | + in_compare, as_constant = _name_mentions(func, name) |
| 128 | + if in_compare: |
| 129 | + return func, cls |
| 130 | + if as_constant and constant_match is None: |
| 131 | + constant_match = (func, cls) |
| 132 | + return constant_match |
| 133 | + |
| 134 | + |
| 135 | +def _iter_functions( |
| 136 | + tree: ast.Module, |
| 137 | +) -> list[tuple[FunctionNode, ast.ClassDef | None]]: |
| 138 | + """ |
| 139 | + Test-file functions paired with their enclosing class (if any) |
| 140 | + """ |
| 141 | + out: list[tuple[FunctionNode, ast.ClassDef | None]] = [] |
| 142 | + for stmt in tree.body: |
| 143 | + if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)): |
| 144 | + out.append((stmt, None)) |
| 145 | + elif isinstance(stmt, ast.ClassDef): |
| 146 | + out.extend( |
| 147 | + (sub, stmt) |
| 148 | + for sub in stmt.body |
| 149 | + if isinstance(sub, (ast.FunctionDef, ast.AsyncFunctionDef)) |
| 150 | + ) |
| 151 | + return out |
| 152 | + |
| 153 | + |
| 154 | +def _name_mentions(func: FunctionNode, name: str) -> tuple[bool, bool]: |
| 155 | + """ |
| 156 | + How `func` mentions `name`: (in a comparison, as any constant) |
| 157 | + """ |
| 158 | + in_compare = False |
| 159 | + as_constant = False |
| 160 | + for node in ast.walk(func): |
| 161 | + if isinstance(node, ast.Constant) and node.value == name: |
| 162 | + as_constant = True |
| 163 | + if isinstance(node, ast.Compare): |
| 164 | + operands = [node.left, *node.comparators] |
| 165 | + if any( |
| 166 | + isinstance(o, ast.Constant) and o.value == name |
| 167 | + for o in operands |
| 168 | + ): |
| 169 | + in_compare = True |
| 170 | + return in_compare, as_constant |
| 171 | + |
| 172 | + |
| 173 | +def _loaded_names(func: FunctionNode) -> set[str]: |
| 174 | + """ |
| 175 | + Names that `func` reads, i.e. its candidate setup dependencies |
| 176 | + """ |
| 177 | + return { |
| 178 | + node.id |
| 179 | + for node in ast.walk(func) |
| 180 | + if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load) |
| 181 | + } |
| 182 | + |
| 183 | + |
| 184 | +def _module_assignments( |
| 185 | + tree: ast.Module, |
| 186 | +) -> list[tuple[set[str], ast.stmt]]: |
| 187 | + """ |
| 188 | + Module-level assignments as (target names, statement) pairs |
| 189 | + """ |
| 190 | + out: list[tuple[set[str], ast.stmt]] = [] |
| 191 | + for stmt in tree.body: |
| 192 | + if isinstance(stmt, ast.Assign): |
| 193 | + targets = {t.id for t in stmt.targets if isinstance(t, ast.Name)} |
| 194 | + elif isinstance(stmt, ast.AnnAssign) and isinstance( |
| 195 | + stmt.target, ast.Name |
| 196 | + ): |
| 197 | + targets = {stmt.target.id} |
| 198 | + else: |
| 199 | + continue |
| 200 | + if targets: |
| 201 | + out.append((targets, stmt)) |
| 202 | + return out |
| 203 | + |
| 204 | + |
| 205 | +def _segment(lines: list[str], node: ast.stmt) -> str: |
| 206 | + """ |
| 207 | + Source text of `node`, including any decorators |
| 208 | + """ |
| 209 | + start = node.lineno |
| 210 | + for deco in getattr(node, "decorator_list", []): |
| 211 | + start = min(start, deco.lineno) |
| 212 | + return "\n".join(lines[start - 1 : node.end_lineno]) |
0 commit comments