rewrite.py 47 KB

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