rewrite.py 47 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216
  1. """Rewrite assertion AST to produce nice error messages."""
  2. from __future__ import annotations
  3. import ast
  4. from collections import defaultdict
  5. from collections.abc import Callable
  6. from collections.abc import Iterable
  7. from collections.abc import Iterator
  8. from collections.abc import Sequence
  9. import errno
  10. import functools
  11. import importlib.abc
  12. import importlib.machinery
  13. import importlib.util
  14. import io
  15. import itertools
  16. import marshal
  17. import os
  18. from pathlib import Path
  19. from pathlib import PurePath
  20. import struct
  21. import sys
  22. import tokenize
  23. import types
  24. from typing import IO
  25. from typing import TYPE_CHECKING
  26. from _pytest._io.saferepr import DEFAULT_REPR_MAX_SIZE
  27. from _pytest._io.saferepr import saferepr
  28. from _pytest._io.saferepr import saferepr_unlimited
  29. from _pytest._version import version
  30. from _pytest.assertion import util
  31. from _pytest.config import Config
  32. from _pytest.fixtures import FixtureFunctionDefinition
  33. from _pytest.main import Session
  34. from _pytest.pathlib import absolutepath
  35. from _pytest.pathlib import fnmatch_ex
  36. from _pytest.stash import StashKey
  37. # fmt: off
  38. from _pytest.assertion.util import format_explanation as _format_explanation # noqa:F401, isort:skip
  39. # fmt:on
  40. if TYPE_CHECKING:
  41. from _pytest.assertion import AssertionState
  42. class Sentinel:
  43. pass
  44. assertstate_key = StashKey["AssertionState"]()
  45. # pytest caches rewritten pycs in pycache dirs
  46. PYTEST_TAG = f"{sys.implementation.cache_tag}-pytest-{version}"
  47. PYC_EXT = ".py" + ((__debug__ and "c") or "o")
  48. PYC_TAIL = "." + PYTEST_TAG + PYC_EXT
  49. # Special marker that denotes we have just left a scope definition
  50. _SCOPE_END_MARKER = Sentinel()
  51. class AssertionRewritingHook(importlib.abc.MetaPathFinder, importlib.abc.Loader):
  52. """PEP302/PEP451 import hook which rewrites asserts."""
  53. def __init__(self, config: Config) -> None:
  54. self.config = config
  55. try:
  56. self.fnpats = config.getini("python_files")
  57. except ValueError:
  58. self.fnpats = ["test_*.py", "*_test.py"]
  59. self.session: Session | None = None
  60. self._rewritten_names: dict[str, Path] = {}
  61. self._must_rewrite: set[str] = set()
  62. # flag to guard against trying to rewrite a pyc file while we are already writing another pyc file,
  63. # which might result in infinite recursion (#3506)
  64. self._writing_pyc = False
  65. self._basenames_to_check_rewrite = {"conftest"}
  66. self._marked_for_rewrite_cache: dict[str, bool] = {}
  67. self._session_paths_checked = False
  68. def set_session(self, session: Session | None) -> None:
  69. self.session = session
  70. self._session_paths_checked = False
  71. # Indirection so we can mock calls to find_spec originated from the hook during testing
  72. _find_spec = importlib.machinery.PathFinder.find_spec
  73. def find_spec(
  74. self,
  75. name: str,
  76. path: Sequence[str | bytes] | None = None,
  77. target: types.ModuleType | None = None,
  78. ) -> importlib.machinery.ModuleSpec | None:
  79. if self._writing_pyc:
  80. return None
  81. state = self.config.stash[assertstate_key]
  82. if self._early_rewrite_bailout(name, state):
  83. return None
  84. state.trace(f"find_module called for: {name}")
  85. # Type ignored because mypy is confused about the `self` binding here.
  86. spec = self._find_spec(name, path) # type: ignore
  87. if spec is None and path is not None:
  88. # With --import-mode=importlib, PathFinder cannot find spec without modifying `sys.path`,
  89. # causing inability to assert rewriting (#12659).
  90. # At this point, try using the file path to find the module spec.
  91. for _path_str in path:
  92. spec = importlib.util.spec_from_file_location(name, _path_str)
  93. if spec is not None:
  94. break
  95. if (
  96. # the import machinery could not find a file to import
  97. spec is None
  98. # this is a namespace package (without `__init__.py`)
  99. # there's nothing to rewrite there
  100. or spec.origin is None
  101. # we can only rewrite source files
  102. or not isinstance(spec.loader, importlib.machinery.SourceFileLoader)
  103. # if the file doesn't exist, we can't rewrite it
  104. or not os.path.exists(spec.origin)
  105. ):
  106. return None
  107. else:
  108. fn = spec.origin
  109. if not self._should_rewrite(name, fn, state):
  110. return None
  111. return importlib.util.spec_from_file_location(
  112. name,
  113. fn,
  114. loader=self,
  115. submodule_search_locations=spec.submodule_search_locations,
  116. )
  117. def create_module(
  118. self, spec: importlib.machinery.ModuleSpec
  119. ) -> types.ModuleType | None:
  120. return None # default behaviour is fine
  121. def exec_module(self, module: types.ModuleType) -> None:
  122. assert module.__spec__ is not None
  123. assert module.__spec__.origin is not None
  124. fn = Path(module.__spec__.origin)
  125. state = self.config.stash[assertstate_key]
  126. self._rewritten_names[module.__name__] = fn
  127. # The requested module looks like a test file, so rewrite it. This is
  128. # the most magical part of the process: load the source, rewrite the
  129. # asserts, and load the rewritten source. We also cache the rewritten
  130. # module code in a special pyc. We must be aware of the possibility of
  131. # concurrent pytest processes rewriting and loading pycs. To avoid
  132. # tricky race conditions, we maintain the following invariant: The
  133. # cached pyc is always a complete, valid pyc. Operations on it must be
  134. # atomic. POSIX's atomic rename comes in handy.
  135. write = not sys.dont_write_bytecode
  136. cache_dir = get_cache_dir(fn)
  137. if write:
  138. ok = try_makedirs(cache_dir)
  139. if not ok:
  140. write = False
  141. state.trace(f"read only directory: {cache_dir}")
  142. cache_name = fn.name[:-3] + PYC_TAIL
  143. pyc = cache_dir / cache_name
  144. # Notice that even if we're in a read-only directory, I'm going
  145. # to check for a cached pyc. This may not be optimal...
  146. co = _read_pyc(fn, pyc, state.trace)
  147. if co is None:
  148. state.trace(f"rewriting {fn!r}")
  149. source_stat, co = _rewrite_test(fn, self.config)
  150. if write:
  151. self._writing_pyc = True
  152. try:
  153. _write_pyc(state, co, source_stat, pyc)
  154. finally:
  155. self._writing_pyc = False
  156. else:
  157. state.trace(f"found cached rewritten pyc for {fn}")
  158. exec(co, module.__dict__)
  159. def _early_rewrite_bailout(self, name: str, state: AssertionState) -> bool:
  160. """A fast way to get out of rewriting modules.
  161. Profiling has shown that the call to PathFinder.find_spec (inside of
  162. the find_spec from this class) is a major slowdown, so, this method
  163. tries to filter what we're sure won't be rewritten before getting to
  164. it.
  165. """
  166. if self.session is not None and not self._session_paths_checked:
  167. self._session_paths_checked = True
  168. for initial_path in self.session._initialpaths:
  169. # Make something as c:/projects/my_project/path.py ->
  170. # ['c:', 'projects', 'my_project', 'path.py']
  171. parts = str(initial_path).split(os.sep)
  172. # add 'path' to basenames to be checked.
  173. self._basenames_to_check_rewrite.add(os.path.splitext(parts[-1])[0])
  174. # Note: conftest already by default in _basenames_to_check_rewrite.
  175. parts = name.split(".")
  176. if parts[-1] in self._basenames_to_check_rewrite:
  177. return False
  178. # For matching the name it must be as if it was a filename.
  179. path = PurePath(*parts).with_suffix(".py")
  180. for pat in self.fnpats:
  181. # if the pattern contains subdirectories ("tests/**.py" for example) we can't bail out based
  182. # on the name alone because we need to match against the full path
  183. if os.path.dirname(pat):
  184. return False
  185. if fnmatch_ex(pat, path):
  186. return False
  187. if self._is_marked_for_rewrite(name, state):
  188. return False
  189. state.trace(f"early skip of rewriting module: {name}")
  190. return True
  191. def _should_rewrite(self, name: str, fn: str, state: AssertionState) -> bool:
  192. # always rewrite conftest files
  193. if os.path.basename(fn) == "conftest.py":
  194. state.trace(f"rewriting conftest file: {fn!r}")
  195. return True
  196. if self.session is not None:
  197. if self.session.isinitpath(absolutepath(fn)):
  198. state.trace(f"matched test file (was specified on cmdline): {fn!r}")
  199. return True
  200. # modules not passed explicitly on the command line are only
  201. # rewritten if they match the naming convention for test files
  202. fn_path = PurePath(fn)
  203. for pat in self.fnpats:
  204. if fnmatch_ex(pat, fn_path):
  205. state.trace(f"matched test file {fn!r}")
  206. return True
  207. return self._is_marked_for_rewrite(name, state)
  208. def _is_marked_for_rewrite(self, name: str, state: AssertionState) -> bool:
  209. try:
  210. return self._marked_for_rewrite_cache[name]
  211. except KeyError:
  212. for marked in self._must_rewrite:
  213. if name == marked or name.startswith(marked + "."):
  214. state.trace(f"matched marked file {name!r} (from {marked!r})")
  215. self._marked_for_rewrite_cache[name] = True
  216. return True
  217. self._marked_for_rewrite_cache[name] = False
  218. return False
  219. def mark_rewrite(self, *names: str) -> None:
  220. """Mark import names as needing to be rewritten.
  221. The named module or package as well as any nested modules will
  222. be rewritten on import.
  223. """
  224. already_imported = (
  225. set(names).intersection(sys.modules).difference(self._rewritten_names)
  226. )
  227. for name in already_imported:
  228. mod = sys.modules[name]
  229. if not AssertionRewriter.is_rewrite_disabled(
  230. mod.__doc__ or ""
  231. ) and not isinstance(mod.__loader__, type(self)):
  232. self._warn_already_imported(name)
  233. self._must_rewrite.update(names)
  234. self._marked_for_rewrite_cache.clear()
  235. def _warn_already_imported(self, name: str) -> None:
  236. from _pytest.warning_types import PytestAssertRewriteWarning
  237. self.config.issue_config_time_warning(
  238. PytestAssertRewriteWarning(
  239. f"Module already imported so cannot be rewritten; {name}"
  240. ),
  241. stacklevel=5,
  242. )
  243. def get_data(self, pathname: str | bytes) -> bytes:
  244. """Optional PEP302 get_data API."""
  245. with open(pathname, "rb") as f:
  246. return f.read()
  247. if sys.version_info >= (3, 10):
  248. if sys.version_info >= (3, 12):
  249. from importlib.resources.abc import TraversableResources
  250. else:
  251. from importlib.abc import TraversableResources
  252. def get_resource_reader(self, name: str) -> TraversableResources:
  253. if sys.version_info < (3, 11):
  254. from importlib.readers import FileReader
  255. else:
  256. from importlib.resources.readers import FileReader
  257. return FileReader(types.SimpleNamespace(path=self._rewritten_names[name]))
  258. def _write_pyc_fp(
  259. fp: IO[bytes], source_stat: os.stat_result, co: types.CodeType
  260. ) -> None:
  261. # Technically, we don't have to have the same pyc format as
  262. # (C)Python, since these "pycs" should never be seen by builtin
  263. # import. However, there's little reason to deviate.
  264. fp.write(importlib.util.MAGIC_NUMBER)
  265. # https://www.python.org/dev/peps/pep-0552/
  266. flags = b"\x00\x00\x00\x00"
  267. fp.write(flags)
  268. # as of now, bytecode header expects 32-bit numbers for size and mtime (#4903)
  269. mtime = int(source_stat.st_mtime) & 0xFFFFFFFF
  270. size = source_stat.st_size & 0xFFFFFFFF
  271. # "<LL" stands for 2 unsigned longs, little-endian.
  272. fp.write(struct.pack("<LL", mtime, size))
  273. fp.write(marshal.dumps(co))
  274. def _write_pyc(
  275. state: AssertionState,
  276. co: types.CodeType,
  277. source_stat: os.stat_result,
  278. pyc: Path,
  279. ) -> bool:
  280. proc_pyc = f"{pyc}.{os.getpid()}"
  281. try:
  282. with open(proc_pyc, "wb") as fp:
  283. _write_pyc_fp(fp, source_stat, co)
  284. except OSError as e:
  285. state.trace(f"error writing pyc file at {proc_pyc}: errno={e.errno}")
  286. return False
  287. try:
  288. os.replace(proc_pyc, pyc)
  289. except OSError as e:
  290. state.trace(f"error writing pyc file at {pyc}: {e}")
  291. # we ignore any failure to write the cache file
  292. # there are many reasons, permission-denied, pycache dir being a
  293. # file etc.
  294. return False
  295. return True
  296. def _rewrite_test(fn: Path, config: Config) -> tuple[os.stat_result, types.CodeType]:
  297. """Read and rewrite *fn* and return the code object."""
  298. stat = os.stat(fn)
  299. source = fn.read_bytes()
  300. strfn = str(fn)
  301. tree = ast.parse(source, filename=strfn)
  302. rewrite_asserts(tree, source, strfn, config)
  303. co = compile(tree, strfn, "exec", dont_inherit=True)
  304. return stat, co
  305. def _read_pyc(
  306. source: Path, pyc: Path, trace: Callable[[str], None] = lambda x: None
  307. ) -> types.CodeType | None:
  308. """Possibly read a pytest pyc containing rewritten code.
  309. Return rewritten code if successful or None if not.
  310. """
  311. try:
  312. fp = open(pyc, "rb")
  313. except OSError:
  314. return None
  315. with fp:
  316. try:
  317. stat_result = os.stat(source)
  318. mtime = int(stat_result.st_mtime)
  319. size = stat_result.st_size
  320. data = fp.read(16)
  321. except OSError as e:
  322. trace(f"_read_pyc({source}): OSError {e}")
  323. return None
  324. # Check for invalid or out of date pyc file.
  325. if len(data) != (16):
  326. trace(f"_read_pyc({source}): invalid pyc (too short)")
  327. return None
  328. if data[:4] != importlib.util.MAGIC_NUMBER:
  329. trace(f"_read_pyc({source}): invalid pyc (bad magic number)")
  330. return None
  331. if data[4:8] != b"\x00\x00\x00\x00":
  332. trace(f"_read_pyc({source}): invalid pyc (unsupported flags)")
  333. return None
  334. mtime_data = data[8:12]
  335. if int.from_bytes(mtime_data, "little") != mtime & 0xFFFFFFFF:
  336. trace(f"_read_pyc({source}): out of date")
  337. return None
  338. size_data = data[12:16]
  339. if int.from_bytes(size_data, "little") != size & 0xFFFFFFFF:
  340. trace(f"_read_pyc({source}): invalid pyc (incorrect size)")
  341. return None
  342. try:
  343. co = marshal.load(fp)
  344. except Exception as e:
  345. trace(f"_read_pyc({source}): marshal.load error {e}")
  346. return None
  347. if not isinstance(co, types.CodeType):
  348. trace(f"_read_pyc({source}): not a code object")
  349. return None
  350. return co
  351. def rewrite_asserts(
  352. mod: ast.Module,
  353. source: bytes,
  354. module_path: str | None = None,
  355. config: Config | None = None,
  356. ) -> None:
  357. """Rewrite the assert statements in mod."""
  358. AssertionRewriter(module_path, config, source).run(mod)
  359. def _saferepr(obj: object) -> str:
  360. r"""Get a safe repr of an object for assertion error messages.
  361. The assertion formatting (util.format_explanation()) requires
  362. newlines to be escaped since they are a special character for it.
  363. Normally assertion.util.format_explanation() does this but for a
  364. custom repr it is possible to contain one of the special escape
  365. sequences, especially '\n{' and '\n}' are likely to be present in
  366. JSON reprs.
  367. """
  368. if isinstance(obj, types.MethodType):
  369. # for bound methods, skip redundant <bound method ...> information
  370. return obj.__name__
  371. maxsize = _get_maxsize_for_saferepr(util._config)
  372. if not maxsize:
  373. return saferepr_unlimited(obj).replace("\n", "\\n")
  374. return saferepr(obj, maxsize=maxsize).replace("\n", "\\n")
  375. def _get_maxsize_for_saferepr(config: Config | None) -> int | None:
  376. """Get `maxsize` configuration for saferepr based on the given config object."""
  377. if config is None:
  378. verbosity = 0
  379. else:
  380. verbosity = config.get_verbosity(Config.VERBOSITY_ASSERTIONS)
  381. if verbosity >= 2:
  382. return None
  383. if verbosity >= 1:
  384. return DEFAULT_REPR_MAX_SIZE * 10
  385. return DEFAULT_REPR_MAX_SIZE
  386. def _format_assertmsg(obj: object) -> str:
  387. r"""Format the custom assertion message given.
  388. For strings this simply replaces newlines with '\n~' so that
  389. util.format_explanation() will preserve them instead of escaping
  390. newlines. For other objects saferepr() is used first.
  391. """
  392. # reprlib appears to have a bug which means that if a string
  393. # contains a newline it gets escaped, however if an object has a
  394. # .__repr__() which contains newlines it does not get escaped.
  395. # However in either case we want to preserve the newline.
  396. replaces = [("\n", "\n~"), ("%", "%%")]
  397. if not isinstance(obj, str):
  398. obj = saferepr(obj, _get_maxsize_for_saferepr(util._config))
  399. replaces.append(("\\n", "\n~"))
  400. for r1, r2 in replaces:
  401. obj = obj.replace(r1, r2)
  402. return obj
  403. def _should_repr_global_name(obj: object) -> bool:
  404. if callable(obj):
  405. # For pytest fixtures the __repr__ method provides more information than the function name.
  406. return isinstance(obj, FixtureFunctionDefinition)
  407. try:
  408. return not hasattr(obj, "__name__")
  409. except Exception:
  410. return True
  411. def _format_boolop(explanations: Iterable[str], is_or: bool) -> str:
  412. explanation = "(" + ((is_or and " or ") or " and ").join(explanations) + ")"
  413. return explanation.replace("%", "%%")
  414. def _call_reprcompare(
  415. ops: Sequence[str],
  416. results: Sequence[bool],
  417. expls: Sequence[str],
  418. each_obj: Sequence[object],
  419. ) -> str:
  420. for i, res, expl in zip(range(len(ops)), results, expls):
  421. try:
  422. done = not res
  423. except Exception:
  424. done = True
  425. if done:
  426. break
  427. if util._reprcompare is not None:
  428. custom = util._reprcompare(ops[i], each_obj[i], each_obj[i + 1])
  429. if custom is not None:
  430. return custom
  431. return expl
  432. def _call_assertion_pass(lineno: int, orig: str, expl: str) -> None:
  433. if util._assertion_pass is not None:
  434. util._assertion_pass(lineno, orig, expl)
  435. def _check_if_assertion_pass_impl() -> bool:
  436. """Check if any plugins implement the pytest_assertion_pass hook
  437. in order not to generate explanation unnecessarily (might be expensive)."""
  438. return True if util._assertion_pass else False
  439. UNARY_MAP = {ast.Not: "not %s", ast.Invert: "~%s", ast.USub: "-%s", ast.UAdd: "+%s"}
  440. BINOP_MAP = {
  441. ast.BitOr: "|",
  442. ast.BitXor: "^",
  443. ast.BitAnd: "&",
  444. ast.LShift: "<<",
  445. ast.RShift: ">>",
  446. ast.Add: "+",
  447. ast.Sub: "-",
  448. ast.Mult: "*",
  449. ast.Div: "/",
  450. ast.FloorDiv: "//",
  451. ast.Mod: "%%", # escaped for string formatting
  452. ast.Eq: "==",
  453. ast.NotEq: "!=",
  454. ast.Lt: "<",
  455. ast.LtE: "<=",
  456. ast.Gt: ">",
  457. ast.GtE: ">=",
  458. ast.Pow: "**",
  459. ast.Is: "is",
  460. ast.IsNot: "is not",
  461. ast.In: "in",
  462. ast.NotIn: "not in",
  463. ast.MatMult: "@",
  464. }
  465. def traverse_node(node: ast.AST) -> Iterator[ast.AST]:
  466. """Recursively yield node and all its children in depth-first order."""
  467. yield node
  468. for child in ast.iter_child_nodes(node):
  469. yield from traverse_node(child)
  470. @functools.lru_cache(maxsize=1)
  471. def _get_assertion_exprs(src: bytes) -> dict[int, str]:
  472. """Return a mapping from {lineno: "assertion test expression"}."""
  473. ret: dict[int, str] = {}
  474. depth = 0
  475. lines: list[str] = []
  476. assert_lineno: int | None = None
  477. seen_lines: set[int] = set()
  478. def _write_and_reset() -> None:
  479. nonlocal depth, lines, assert_lineno, seen_lines
  480. assert assert_lineno is not None
  481. ret[assert_lineno] = "".join(lines).rstrip().rstrip("\\")
  482. depth = 0
  483. lines = []
  484. assert_lineno = None
  485. seen_lines = set()
  486. tokens = tokenize.tokenize(io.BytesIO(src).readline)
  487. for tp, source, (lineno, offset), _, line in tokens:
  488. if tp == tokenize.NAME and source == "assert":
  489. assert_lineno = lineno
  490. elif assert_lineno is not None:
  491. # keep track of depth for the assert-message `,` lookup
  492. if tp == tokenize.OP and source in "([{":
  493. depth += 1
  494. elif tp == tokenize.OP and source in ")]}":
  495. depth -= 1
  496. if not lines:
  497. lines.append(line[offset:])
  498. seen_lines.add(lineno)
  499. # a non-nested comma separates the expression from the message
  500. elif depth == 0 and tp == tokenize.OP and source == ",":
  501. # one line assert with message
  502. if lineno in seen_lines and len(lines) == 1:
  503. offset_in_trimmed = offset + len(lines[-1]) - len(line)
  504. lines[-1] = lines[-1][:offset_in_trimmed]
  505. # multi-line assert with message
  506. elif lineno in seen_lines:
  507. lines[-1] = lines[-1][:offset]
  508. # multi line assert with escaped newline before message
  509. else:
  510. lines.append(line[:offset])
  511. _write_and_reset()
  512. elif tp in {tokenize.NEWLINE, tokenize.ENDMARKER}:
  513. _write_and_reset()
  514. elif lines and lineno not in seen_lines:
  515. lines.append(line)
  516. seen_lines.add(lineno)
  517. return ret
  518. class AssertionRewriter(ast.NodeVisitor):
  519. """Assertion rewriting implementation.
  520. The main entrypoint is to call .run() with an ast.Module instance,
  521. this will then find all the assert statements and rewrite them to
  522. provide intermediate values and a detailed assertion error. See
  523. http://pybites.blogspot.be/2011/07/behind-scenes-of-pytests-new-assertion.html
  524. for an overview of how this works.
  525. The entry point here is .run() which will iterate over all the
  526. statements in an ast.Module and for each ast.Assert statement it
  527. finds call .visit() with it. Then .visit_Assert() takes over and
  528. is responsible for creating new ast statements to replace the
  529. original assert statement: it rewrites the test of an assertion
  530. to provide intermediate values and replace it with an if statement
  531. which raises an assertion error with a detailed explanation in
  532. case the expression is false and calls pytest_assertion_pass hook
  533. if expression is true.
  534. For this .visit_Assert() uses the visitor pattern to visit all the
  535. AST nodes of the ast.Assert.test field, each visit call returning
  536. an AST node and the corresponding explanation string. During this
  537. state is kept in several instance attributes:
  538. :statements: All the AST statements which will replace the assert
  539. statement.
  540. :variables: This is populated by .variable() with each variable
  541. used by the statements so that they can all be set to None at
  542. the end of the statements.
  543. :variable_counter: Counter to create new unique variables needed
  544. by statements. Variables are created using .variable() and
  545. have the form of "@py_assert0".
  546. :expl_stmts: The AST statements which will be executed to get
  547. data from the assertion. This is the code which will construct
  548. the detailed assertion message that is used in the AssertionError
  549. or for the pytest_assertion_pass hook.
  550. :explanation_specifiers: A dict filled by .explanation_param()
  551. with %-formatting placeholders and their corresponding
  552. expressions to use in the building of an assertion message.
  553. This is used by .pop_format_context() to build a message.
  554. :stack: A stack of the explanation_specifiers dicts maintained by
  555. .push_format_context() and .pop_format_context() which allows
  556. to build another %-formatted string while already building one.
  557. :scope: A tuple containing the current scope used for variables_overwrite.
  558. :variables_overwrite: A dict filled with references to variables
  559. that change value within an assert. This happens when a variable is
  560. reassigned with the walrus operator
  561. This state, except the variables_overwrite, is reset on every new assert
  562. statement visited and used by the other visitors.
  563. """
  564. def __init__(
  565. self, module_path: str | None, config: Config | None, source: bytes
  566. ) -> None:
  567. super().__init__()
  568. self.module_path = module_path
  569. self.config = config
  570. if config is not None:
  571. self.enable_assertion_pass_hook = config.getini(
  572. "enable_assertion_pass_hook"
  573. )
  574. else:
  575. self.enable_assertion_pass_hook = False
  576. self.source = source
  577. self.scope: tuple[ast.AST, ...] = ()
  578. self.variables_overwrite: defaultdict[tuple[ast.AST, ...], dict[str, str]] = (
  579. defaultdict(dict)
  580. )
  581. def run(self, mod: ast.Module) -> None:
  582. """Find all assert statements in *mod* and rewrite them."""
  583. if not mod.body:
  584. # Nothing to do.
  585. return
  586. # We'll insert some special imports at the top of the module, but after any
  587. # docstrings and __future__ imports, so first figure out where that is.
  588. doc = getattr(mod, "docstring", None)
  589. expect_docstring = doc is None
  590. if doc is not None and self.is_rewrite_disabled(doc):
  591. return
  592. pos = 0
  593. item = None
  594. for item in mod.body:
  595. if (
  596. expect_docstring
  597. and isinstance(item, ast.Expr)
  598. and isinstance(item.value, ast.Constant)
  599. and isinstance(item.value.value, str)
  600. ):
  601. doc = item.value.value
  602. if self.is_rewrite_disabled(doc):
  603. return
  604. expect_docstring = False
  605. elif (
  606. isinstance(item, ast.ImportFrom)
  607. and item.level == 0
  608. and item.module == "__future__"
  609. ):
  610. pass
  611. else:
  612. break
  613. pos += 1
  614. # Special case: for a decorated function, set the lineno to that of the
  615. # first decorator, not the `def`. Issue #4984.
  616. if isinstance(item, ast.FunctionDef) and item.decorator_list:
  617. lineno = item.decorator_list[0].lineno
  618. else:
  619. lineno = item.lineno
  620. # Now actually insert the special imports.
  621. if sys.version_info >= (3, 10):
  622. aliases = [
  623. ast.alias("builtins", "@py_builtins", lineno=lineno, col_offset=0),
  624. ast.alias(
  625. "_pytest.assertion.rewrite",
  626. "@pytest_ar",
  627. lineno=lineno,
  628. col_offset=0,
  629. ),
  630. ]
  631. else:
  632. aliases = [
  633. ast.alias("builtins", "@py_builtins"),
  634. ast.alias("_pytest.assertion.rewrite", "@pytest_ar"),
  635. ]
  636. imports = [
  637. ast.Import([alias], lineno=lineno, col_offset=0) for alias in aliases
  638. ]
  639. mod.body[pos:pos] = imports
  640. # Collect asserts.
  641. self.scope = (mod,)
  642. nodes: list[ast.AST | Sentinel] = [mod]
  643. while nodes:
  644. node = nodes.pop()
  645. if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
  646. self.scope = tuple((*self.scope, node))
  647. nodes.append(_SCOPE_END_MARKER)
  648. if node == _SCOPE_END_MARKER:
  649. self.scope = self.scope[:-1]
  650. continue
  651. assert isinstance(node, ast.AST)
  652. for name, field in ast.iter_fields(node):
  653. if isinstance(field, list):
  654. new: list[ast.AST] = []
  655. for i, child in enumerate(field):
  656. if isinstance(child, ast.Assert):
  657. # Transform assert.
  658. new.extend(self.visit(child))
  659. else:
  660. new.append(child)
  661. if isinstance(child, ast.AST):
  662. nodes.append(child)
  663. setattr(node, name, new)
  664. elif (
  665. isinstance(field, ast.AST)
  666. # Don't recurse into expressions as they can't contain
  667. # asserts.
  668. and not isinstance(field, ast.expr)
  669. ):
  670. nodes.append(field)
  671. @staticmethod
  672. def is_rewrite_disabled(docstring: str) -> bool:
  673. return "PYTEST_DONT_REWRITE" in docstring
  674. def variable(self) -> str:
  675. """Get a new variable."""
  676. # Use a character invalid in python identifiers to avoid clashing.
  677. name = "@py_assert" + str(next(self.variable_counter))
  678. self.variables.append(name)
  679. return name
  680. def assign(self, expr: ast.expr) -> ast.Name:
  681. """Give *expr* a name."""
  682. name = self.variable()
  683. self.statements.append(ast.Assign([ast.Name(name, ast.Store())], expr))
  684. return ast.copy_location(ast.Name(name, ast.Load()), expr)
  685. def display(self, expr: ast.expr) -> ast.expr:
  686. """Call saferepr on the expression."""
  687. return self.helper("_saferepr", expr)
  688. def helper(self, name: str, *args: ast.expr) -> ast.expr:
  689. """Call a helper in this module."""
  690. py_name = ast.Name("@pytest_ar", ast.Load())
  691. attr = ast.Attribute(py_name, name, ast.Load())
  692. return ast.Call(attr, list(args), [])
  693. def builtin(self, name: str) -> ast.Attribute:
  694. """Return the builtin called *name*."""
  695. builtin_name = ast.Name("@py_builtins", ast.Load())
  696. return ast.Attribute(builtin_name, name, ast.Load())
  697. def explanation_param(self, expr: ast.expr) -> str:
  698. """Return a new named %-formatting placeholder for expr.
  699. This creates a %-formatting placeholder for expr in the
  700. current formatting context, e.g. ``%(py0)s``. The placeholder
  701. and expr are placed in the current format context so that it
  702. can be used on the next call to .pop_format_context().
  703. """
  704. specifier = "py" + str(next(self.variable_counter))
  705. self.explanation_specifiers[specifier] = expr
  706. return "%(" + specifier + ")s"
  707. def push_format_context(self) -> None:
  708. """Create a new formatting context.
  709. The format context is used for when an explanation wants to
  710. have a variable value formatted in the assertion message. In
  711. this case the value required can be added using
  712. .explanation_param(). Finally .pop_format_context() is used
  713. to format a string of %-formatted values as added by
  714. .explanation_param().
  715. """
  716. self.explanation_specifiers: dict[str, ast.expr] = {}
  717. self.stack.append(self.explanation_specifiers)
  718. def pop_format_context(self, expl_expr: ast.expr) -> ast.Name:
  719. """Format the %-formatted string with current format context.
  720. The expl_expr should be an str ast.expr instance constructed from
  721. the %-placeholders created by .explanation_param(). This will
  722. add the required code to format said string to .expl_stmts and
  723. return the ast.Name instance of the formatted string.
  724. """
  725. current = self.stack.pop()
  726. if self.stack:
  727. self.explanation_specifiers = self.stack[-1]
  728. keys: list[ast.expr | None] = [ast.Constant(key) for key in current.keys()]
  729. format_dict = ast.Dict(keys, list(current.values()))
  730. form = ast.BinOp(expl_expr, ast.Mod(), format_dict)
  731. name = "@py_format" + str(next(self.variable_counter))
  732. if self.enable_assertion_pass_hook:
  733. self.format_variables.append(name)
  734. self.expl_stmts.append(ast.Assign([ast.Name(name, ast.Store())], form))
  735. return ast.Name(name, ast.Load())
  736. def generic_visit(self, node: ast.AST) -> tuple[ast.Name, str]:
  737. """Handle expressions we don't have custom code for."""
  738. assert isinstance(node, ast.expr)
  739. res = self.assign(node)
  740. return res, self.explanation_param(self.display(res))
  741. def visit_Assert(self, assert_: ast.Assert) -> list[ast.stmt]:
  742. """Return the AST statements to replace the ast.Assert instance.
  743. This rewrites the test of an assertion to provide
  744. intermediate values and replace it with an if statement which
  745. raises an assertion error with a detailed explanation in case
  746. the expression is false.
  747. """
  748. if isinstance(assert_.test, ast.Tuple) and len(assert_.test.elts) >= 1:
  749. import warnings
  750. from _pytest.warning_types import PytestAssertRewriteWarning
  751. # TODO: This assert should not be needed.
  752. assert self.module_path is not None
  753. warnings.warn_explicit(
  754. PytestAssertRewriteWarning(
  755. "assertion is always true, perhaps remove parentheses?"
  756. ),
  757. category=None,
  758. filename=self.module_path,
  759. lineno=assert_.lineno,
  760. )
  761. self.statements: list[ast.stmt] = []
  762. self.variables: list[str] = []
  763. self.variable_counter = itertools.count()
  764. if self.enable_assertion_pass_hook:
  765. self.format_variables: list[str] = []
  766. self.stack: list[dict[str, ast.expr]] = []
  767. self.expl_stmts: list[ast.stmt] = []
  768. self.push_format_context()
  769. # Rewrite assert into a bunch of statements.
  770. top_condition, explanation = self.visit(assert_.test)
  771. negation = ast.UnaryOp(ast.Not(), top_condition)
  772. if self.enable_assertion_pass_hook: # Experimental pytest_assertion_pass hook
  773. msg = self.pop_format_context(ast.Constant(explanation))
  774. # Failed
  775. if assert_.msg:
  776. assertmsg = self.helper("_format_assertmsg", assert_.msg)
  777. gluestr = "\n>assert "
  778. else:
  779. assertmsg = ast.Constant("")
  780. gluestr = "assert "
  781. err_explanation = ast.BinOp(ast.Constant(gluestr), ast.Add(), msg)
  782. err_msg = ast.BinOp(assertmsg, ast.Add(), err_explanation)
  783. err_name = ast.Name("AssertionError", ast.Load())
  784. fmt = self.helper("_format_explanation", err_msg)
  785. exc = ast.Call(err_name, [fmt], [])
  786. raise_ = ast.Raise(exc, None)
  787. statements_fail = []
  788. statements_fail.extend(self.expl_stmts)
  789. statements_fail.append(raise_)
  790. # Passed
  791. fmt_pass = self.helper("_format_explanation", msg)
  792. orig = _get_assertion_exprs(self.source)[assert_.lineno]
  793. hook_call_pass = ast.Expr(
  794. self.helper(
  795. "_call_assertion_pass",
  796. ast.Constant(assert_.lineno),
  797. ast.Constant(orig),
  798. fmt_pass,
  799. )
  800. )
  801. # If any hooks implement assert_pass hook
  802. hook_impl_test = ast.If(
  803. self.helper("_check_if_assertion_pass_impl"),
  804. [*self.expl_stmts, hook_call_pass],
  805. [],
  806. )
  807. statements_pass: list[ast.stmt] = [hook_impl_test]
  808. # Test for assertion condition
  809. main_test = ast.If(negation, statements_fail, statements_pass)
  810. self.statements.append(main_test)
  811. if self.format_variables:
  812. variables: list[ast.expr] = [
  813. ast.Name(name, ast.Store()) for name in self.format_variables
  814. ]
  815. clear_format = ast.Assign(variables, ast.Constant(None))
  816. self.statements.append(clear_format)
  817. else: # Original assertion rewriting
  818. # Create failure message.
  819. body = self.expl_stmts
  820. self.statements.append(ast.If(negation, body, []))
  821. if assert_.msg:
  822. assertmsg = self.helper("_format_assertmsg", assert_.msg)
  823. explanation = "\n>assert " + explanation
  824. else:
  825. assertmsg = ast.Constant("")
  826. explanation = "assert " + explanation
  827. template = ast.BinOp(assertmsg, ast.Add(), ast.Constant(explanation))
  828. msg = self.pop_format_context(template)
  829. fmt = self.helper("_format_explanation", msg)
  830. err_name = ast.Name("AssertionError", ast.Load())
  831. exc = ast.Call(err_name, [fmt], [])
  832. raise_ = ast.Raise(exc, None)
  833. body.append(raise_)
  834. # Clear temporary variables by setting them to None.
  835. if self.variables:
  836. variables = [ast.Name(name, ast.Store()) for name in self.variables]
  837. clear = ast.Assign(variables, ast.Constant(None))
  838. self.statements.append(clear)
  839. # Fix locations (line numbers/column offsets).
  840. for stmt in self.statements:
  841. for node in traverse_node(stmt):
  842. if getattr(node, "lineno", None) is None:
  843. # apply the assertion location to all generated ast nodes without source location
  844. # and preserve the location of existing nodes or generated nodes with an correct location.
  845. ast.copy_location(node, assert_)
  846. return self.statements
  847. def visit_NamedExpr(self, name: ast.NamedExpr) -> tuple[ast.NamedExpr, str]:
  848. # This method handles the 'walrus operator' repr of the target
  849. # name if it's a local variable or _should_repr_global_name()
  850. # thinks it's acceptable.
  851. locs = ast.Call(self.builtin("locals"), [], [])
  852. target_id = name.target.id
  853. inlocs = ast.Compare(ast.Constant(target_id), [ast.In()], [locs])
  854. dorepr = self.helper("_should_repr_global_name", name)
  855. test = ast.BoolOp(ast.Or(), [inlocs, dorepr])
  856. expr = ast.IfExp(test, self.display(name), ast.Constant(target_id))
  857. return name, self.explanation_param(expr)
  858. def visit_Name(self, name: ast.Name) -> tuple[ast.Name, str]:
  859. # Display the repr of the name if it's a local variable or
  860. # _should_repr_global_name() thinks it's acceptable.
  861. locs = ast.Call(self.builtin("locals"), [], [])
  862. inlocs = ast.Compare(ast.Constant(name.id), [ast.In()], [locs])
  863. dorepr = self.helper("_should_repr_global_name", name)
  864. test = ast.BoolOp(ast.Or(), [inlocs, dorepr])
  865. expr = ast.IfExp(test, self.display(name), ast.Constant(name.id))
  866. return name, self.explanation_param(expr)
  867. def visit_BoolOp(self, boolop: ast.BoolOp) -> tuple[ast.Name, str]:
  868. res_var = self.variable()
  869. expl_list = self.assign(ast.List([], ast.Load()))
  870. app = ast.Attribute(expl_list, "append", ast.Load())
  871. is_or = int(isinstance(boolop.op, ast.Or))
  872. body = save = self.statements
  873. fail_save = self.expl_stmts
  874. levels = len(boolop.values) - 1
  875. self.push_format_context()
  876. # Process each operand, short-circuiting if needed.
  877. for i, v in enumerate(boolop.values):
  878. if i:
  879. fail_inner: list[ast.stmt] = []
  880. # cond is set in a prior loop iteration below
  881. self.expl_stmts.append(ast.If(cond, fail_inner, [])) # noqa: F821
  882. self.expl_stmts = fail_inner
  883. # Check if the left operand is a ast.NamedExpr and the value has already been visited
  884. if (
  885. isinstance(v, ast.Compare)
  886. and isinstance(v.left, ast.NamedExpr)
  887. and v.left.target.id
  888. in [
  889. ast_expr.id
  890. for ast_expr in boolop.values[:i]
  891. if hasattr(ast_expr, "id")
  892. ]
  893. ):
  894. pytest_temp = self.variable()
  895. self.variables_overwrite[self.scope][v.left.target.id] = v.left # type:ignore[assignment]
  896. v.left.target.id = pytest_temp
  897. self.push_format_context()
  898. res, expl = self.visit(v)
  899. body.append(ast.Assign([ast.Name(res_var, ast.Store())], res))
  900. expl_format = self.pop_format_context(ast.Constant(expl))
  901. call = ast.Call(app, [expl_format], [])
  902. self.expl_stmts.append(ast.Expr(call))
  903. if i < levels:
  904. cond: ast.expr = res
  905. if is_or:
  906. cond = ast.UnaryOp(ast.Not(), cond)
  907. inner: list[ast.stmt] = []
  908. self.statements.append(ast.If(cond, inner, []))
  909. self.statements = body = inner
  910. self.statements = save
  911. self.expl_stmts = fail_save
  912. expl_template = self.helper("_format_boolop", expl_list, ast.Constant(is_or))
  913. expl = self.pop_format_context(expl_template)
  914. return ast.Name(res_var, ast.Load()), self.explanation_param(expl)
  915. def visit_UnaryOp(self, unary: ast.UnaryOp) -> tuple[ast.Name, str]:
  916. pattern = UNARY_MAP[unary.op.__class__]
  917. operand_res, operand_expl = self.visit(unary.operand)
  918. res = self.assign(ast.copy_location(ast.UnaryOp(unary.op, operand_res), unary))
  919. return res, pattern % (operand_expl,)
  920. def visit_BinOp(self, binop: ast.BinOp) -> tuple[ast.Name, str]:
  921. symbol = BINOP_MAP[binop.op.__class__]
  922. left_expr, left_expl = self.visit(binop.left)
  923. right_expr, right_expl = self.visit(binop.right)
  924. explanation = f"({left_expl} {symbol} {right_expl})"
  925. res = self.assign(
  926. ast.copy_location(ast.BinOp(left_expr, binop.op, right_expr), binop)
  927. )
  928. return res, explanation
  929. def visit_Call(self, call: ast.Call) -> tuple[ast.Name, str]:
  930. new_func, func_expl = self.visit(call.func)
  931. arg_expls = []
  932. new_args = []
  933. new_kwargs = []
  934. for arg in call.args:
  935. if isinstance(arg, ast.Name) and arg.id in self.variables_overwrite.get(
  936. self.scope, {}
  937. ):
  938. arg = self.variables_overwrite[self.scope][arg.id] # type:ignore[assignment]
  939. res, expl = self.visit(arg)
  940. arg_expls.append(expl)
  941. new_args.append(res)
  942. for keyword in call.keywords:
  943. if isinstance(
  944. keyword.value, ast.Name
  945. ) and keyword.value.id in self.variables_overwrite.get(self.scope, {}):
  946. keyword.value = self.variables_overwrite[self.scope][keyword.value.id] # type:ignore[assignment]
  947. res, expl = self.visit(keyword.value)
  948. new_kwargs.append(ast.keyword(keyword.arg, res))
  949. if keyword.arg:
  950. arg_expls.append(keyword.arg + "=" + expl)
  951. else: # **args have `arg` keywords with an .arg of None
  952. arg_expls.append("**" + expl)
  953. expl = "{}({})".format(func_expl, ", ".join(arg_expls))
  954. new_call = ast.copy_location(ast.Call(new_func, new_args, new_kwargs), call)
  955. res = self.assign(new_call)
  956. res_expl = self.explanation_param(self.display(res))
  957. outer_expl = f"{res_expl}\n{{{res_expl} = {expl}\n}}"
  958. return res, outer_expl
  959. def visit_Starred(self, starred: ast.Starred) -> tuple[ast.Starred, str]:
  960. # A Starred node can appear in a function call.
  961. res, expl = self.visit(starred.value)
  962. new_starred = ast.Starred(res, starred.ctx)
  963. return new_starred, "*" + expl
  964. def visit_Attribute(self, attr: ast.Attribute) -> tuple[ast.Name, str]:
  965. if not isinstance(attr.ctx, ast.Load):
  966. return self.generic_visit(attr)
  967. value, value_expl = self.visit(attr.value)
  968. res = self.assign(
  969. ast.copy_location(ast.Attribute(value, attr.attr, ast.Load()), attr)
  970. )
  971. res_expl = self.explanation_param(self.display(res))
  972. pat = "%s\n{%s = %s.%s\n}"
  973. expl = pat % (res_expl, res_expl, value_expl, attr.attr)
  974. return res, expl
  975. def visit_Compare(self, comp: ast.Compare) -> tuple[ast.expr, str]:
  976. self.push_format_context()
  977. # We first check if we have overwritten a variable in the previous assert
  978. if isinstance(
  979. comp.left, ast.Name
  980. ) and comp.left.id in self.variables_overwrite.get(self.scope, {}):
  981. comp.left = self.variables_overwrite[self.scope][comp.left.id] # type:ignore[assignment]
  982. if isinstance(comp.left, ast.NamedExpr):
  983. self.variables_overwrite[self.scope][comp.left.target.id] = comp.left # type:ignore[assignment]
  984. left_res, left_expl = self.visit(comp.left)
  985. if isinstance(comp.left, (ast.Compare, ast.BoolOp)):
  986. left_expl = f"({left_expl})"
  987. res_variables = [self.variable() for i in range(len(comp.ops))]
  988. load_names: list[ast.expr] = [ast.Name(v, ast.Load()) for v in res_variables]
  989. store_names = [ast.Name(v, ast.Store()) for v in res_variables]
  990. it = zip(range(len(comp.ops)), comp.ops, comp.comparators)
  991. expls: list[ast.expr] = []
  992. syms: list[ast.expr] = []
  993. results = [left_res]
  994. for i, op, next_operand in it:
  995. if (
  996. isinstance(next_operand, ast.NamedExpr)
  997. and isinstance(left_res, ast.Name)
  998. and next_operand.target.id == left_res.id
  999. ):
  1000. next_operand.target.id = self.variable()
  1001. self.variables_overwrite[self.scope][left_res.id] = next_operand # type:ignore[assignment]
  1002. next_res, next_expl = self.visit(next_operand)
  1003. if isinstance(next_operand, (ast.Compare, ast.BoolOp)):
  1004. next_expl = f"({next_expl})"
  1005. results.append(next_res)
  1006. sym = BINOP_MAP[op.__class__]
  1007. syms.append(ast.Constant(sym))
  1008. expl = f"{left_expl} {sym} {next_expl}"
  1009. expls.append(ast.Constant(expl))
  1010. res_expr = ast.copy_location(ast.Compare(left_res, [op], [next_res]), comp)
  1011. self.statements.append(ast.Assign([store_names[i]], res_expr))
  1012. left_res, left_expl = next_res, next_expl
  1013. # Use pytest.assertion.util._reprcompare if that's available.
  1014. expl_call = self.helper(
  1015. "_call_reprcompare",
  1016. ast.Tuple(syms, ast.Load()),
  1017. ast.Tuple(load_names, ast.Load()),
  1018. ast.Tuple(expls, ast.Load()),
  1019. ast.Tuple(results, ast.Load()),
  1020. )
  1021. if len(comp.ops) > 1:
  1022. res: ast.expr = ast.BoolOp(ast.And(), load_names)
  1023. else:
  1024. res = load_names[0]
  1025. return res, self.explanation_param(self.pop_format_context(expl_call))
  1026. def try_makedirs(cache_dir: Path) -> bool:
  1027. """Attempt to create the given directory and sub-directories exist.
  1028. Returns True if successful or if it already exists.
  1029. """
  1030. try:
  1031. os.makedirs(cache_dir, exist_ok=True)
  1032. except (FileNotFoundError, NotADirectoryError, FileExistsError):
  1033. # One of the path components was not a directory:
  1034. # - we're in a zip file
  1035. # - it is a file
  1036. return False
  1037. except PermissionError:
  1038. return False
  1039. except OSError as e:
  1040. # as of now, EROFS doesn't have an equivalent OSError-subclass
  1041. #
  1042. # squashfuse_ll returns ENOSYS "OSError: [Errno 38] Function not
  1043. # implemented" for a read-only error
  1044. if e.errno in {errno.EROFS, errno.ENOSYS}:
  1045. return False
  1046. raise
  1047. return True
  1048. def get_cache_dir(file_path: Path) -> Path:
  1049. """Return the cache directory to write .pyc files for the given .py file path."""
  1050. if sys.pycache_prefix:
  1051. # given:
  1052. # prefix = '/tmp/pycs'
  1053. # path = '/home/user/proj/test_app.py'
  1054. # we want:
  1055. # '/tmp/pycs/home/user/proj'
  1056. return Path(sys.pycache_prefix) / Path(*file_path.parts[1:-1])
  1057. else:
  1058. # classic pycache directory
  1059. return file_path.parent / "__pycache__"