tmpdir.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315
  1. # mypy: allow-untyped-defs
  2. """Support for providing temporary directories to test functions."""
  3. from __future__ import annotations
  4. from collections.abc import Generator
  5. import dataclasses
  6. import os
  7. from pathlib import Path
  8. import re
  9. from shutil import rmtree
  10. import tempfile
  11. from typing import Any
  12. from typing import final
  13. from typing import Literal
  14. from .pathlib import cleanup_dead_symlinks
  15. from .pathlib import LOCK_TIMEOUT
  16. from .pathlib import make_numbered_dir
  17. from .pathlib import make_numbered_dir_with_cleanup
  18. from .pathlib import rm_rf
  19. from _pytest.compat import get_user_id
  20. from _pytest.config import Config
  21. from _pytest.config import ExitCode
  22. from _pytest.config import hookimpl
  23. from _pytest.config.argparsing import Parser
  24. from _pytest.deprecated import check_ispytest
  25. from _pytest.fixtures import fixture
  26. from _pytest.fixtures import FixtureRequest
  27. from _pytest.monkeypatch import MonkeyPatch
  28. from _pytest.nodes import Item
  29. from _pytest.reports import TestReport
  30. from _pytest.stash import StashKey
  31. tmppath_result_key = StashKey[dict[str, bool]]()
  32. RetentionType = Literal["all", "failed", "none"]
  33. @final
  34. @dataclasses.dataclass
  35. class TempPathFactory:
  36. """Factory for temporary directories under the common base temp directory,
  37. as discussed at :ref:`temporary directory location and retention`.
  38. """
  39. _given_basetemp: Path | None
  40. # pluggy TagTracerSub, not currently exposed, so Any.
  41. _trace: Any
  42. _basetemp: Path | None
  43. _retention_count: int
  44. _retention_policy: RetentionType
  45. def __init__(
  46. self,
  47. given_basetemp: Path | None,
  48. retention_count: int,
  49. retention_policy: RetentionType,
  50. trace,
  51. basetemp: Path | None = None,
  52. *,
  53. _ispytest: bool = False,
  54. ) -> None:
  55. check_ispytest(_ispytest)
  56. if given_basetemp is None:
  57. self._given_basetemp = None
  58. else:
  59. # Use os.path.abspath() to get absolute path instead of resolve() as it
  60. # does not work the same in all platforms (see #4427).
  61. # Path.absolute() exists, but it is not public (see https://bugs.python.org/issue25012).
  62. self._given_basetemp = Path(os.path.abspath(str(given_basetemp)))
  63. self._trace = trace
  64. self._retention_count = retention_count
  65. self._retention_policy = retention_policy
  66. self._basetemp = basetemp
  67. @classmethod
  68. def from_config(
  69. cls,
  70. config: Config,
  71. *,
  72. _ispytest: bool = False,
  73. ) -> TempPathFactory:
  74. """Create a factory according to pytest configuration.
  75. :meta private:
  76. """
  77. check_ispytest(_ispytest)
  78. count = int(config.getini("tmp_path_retention_count"))
  79. if count < 0:
  80. raise ValueError(
  81. f"tmp_path_retention_count must be >= 0. Current input: {count}."
  82. )
  83. policy = config.getini("tmp_path_retention_policy")
  84. if policy not in ("all", "failed", "none"):
  85. raise ValueError(
  86. f"tmp_path_retention_policy must be either all, failed, none. Current input: {policy}."
  87. )
  88. return cls(
  89. given_basetemp=config.option.basetemp,
  90. trace=config.trace.get("tmpdir"),
  91. retention_count=count,
  92. retention_policy=policy,
  93. _ispytest=True,
  94. )
  95. def _ensure_relative_to_basetemp(self, basename: str) -> str:
  96. basename = os.path.normpath(basename)
  97. if (self.getbasetemp() / basename).resolve().parent != self.getbasetemp():
  98. raise ValueError(f"{basename} is not a normalized and relative path")
  99. return basename
  100. def mktemp(self, basename: str, numbered: bool = True) -> Path:
  101. """Create a new temporary directory managed by the factory.
  102. :param basename:
  103. Directory base name, must be a relative path.
  104. :param numbered:
  105. If ``True``, ensure the directory is unique by adding a numbered
  106. suffix greater than any existing one: ``basename="foo-"`` and ``numbered=True``
  107. means that this function will create directories named ``"foo-0"``,
  108. ``"foo-1"``, ``"foo-2"`` and so on.
  109. :returns:
  110. The path to the new directory.
  111. """
  112. basename = self._ensure_relative_to_basetemp(basename)
  113. if not numbered:
  114. p = self.getbasetemp().joinpath(basename)
  115. p.mkdir(mode=0o700)
  116. else:
  117. p = make_numbered_dir(root=self.getbasetemp(), prefix=basename, mode=0o700)
  118. self._trace("mktemp", p)
  119. return p
  120. def getbasetemp(self) -> Path:
  121. """Return the base temporary directory, creating it if needed.
  122. :returns:
  123. The base temporary directory.
  124. """
  125. if self._basetemp is not None:
  126. return self._basetemp
  127. if self._given_basetemp is not None:
  128. basetemp = self._given_basetemp
  129. if basetemp.exists():
  130. rm_rf(basetemp)
  131. basetemp.mkdir(mode=0o700)
  132. basetemp = basetemp.resolve()
  133. else:
  134. from_env = os.environ.get("PYTEST_DEBUG_TEMPROOT")
  135. temproot = Path(from_env or tempfile.gettempdir()).resolve()
  136. user = get_user() or "unknown"
  137. # use a sub-directory in the temproot to speed-up
  138. # make_numbered_dir() call
  139. rootdir = temproot.joinpath(f"pytest-of-{user}")
  140. try:
  141. rootdir.mkdir(mode=0o700, exist_ok=True)
  142. except OSError:
  143. # getuser() likely returned illegal characters for the platform, use unknown back off mechanism
  144. rootdir = temproot.joinpath("pytest-of-unknown")
  145. rootdir.mkdir(mode=0o700, exist_ok=True)
  146. # Because we use exist_ok=True with a predictable name, make sure
  147. # we are the owners, to prevent any funny business (on unix, where
  148. # temproot is usually shared).
  149. # Also, to keep things private, fixup any world-readable temp
  150. # rootdir's permissions. Historically 0o755 was used, so we can't
  151. # just error out on this, at least for a while.
  152. uid = get_user_id()
  153. if uid is not None:
  154. rootdir_stat = rootdir.stat()
  155. if rootdir_stat.st_uid != uid:
  156. raise OSError(
  157. f"The temporary directory {rootdir} is not owned by the current user. "
  158. "Fix this and try again."
  159. )
  160. if (rootdir_stat.st_mode & 0o077) != 0:
  161. os.chmod(rootdir, rootdir_stat.st_mode & ~0o077)
  162. keep = self._retention_count
  163. if self._retention_policy == "none":
  164. keep = 0
  165. basetemp = make_numbered_dir_with_cleanup(
  166. prefix="pytest-",
  167. root=rootdir,
  168. keep=keep,
  169. lock_timeout=LOCK_TIMEOUT,
  170. mode=0o700,
  171. )
  172. assert basetemp is not None, basetemp
  173. self._basetemp = basetemp
  174. self._trace("new basetemp", basetemp)
  175. return basetemp
  176. def get_user() -> str | None:
  177. """Return the current user name, or None if getuser() does not work
  178. in the current environment (see #1010)."""
  179. try:
  180. # In some exotic environments, getpass may not be importable.
  181. import getpass
  182. return getpass.getuser()
  183. except (ImportError, OSError, KeyError):
  184. return None
  185. def pytest_configure(config: Config) -> None:
  186. """Create a TempPathFactory and attach it to the config object.
  187. This is to comply with existing plugins which expect the handler to be
  188. available at pytest_configure time, but ideally should be moved entirely
  189. to the tmp_path_factory session fixture.
  190. """
  191. mp = MonkeyPatch()
  192. config.add_cleanup(mp.undo)
  193. _tmp_path_factory = TempPathFactory.from_config(config, _ispytest=True)
  194. mp.setattr(config, "_tmp_path_factory", _tmp_path_factory, raising=False)
  195. def pytest_addoption(parser: Parser) -> None:
  196. parser.addini(
  197. "tmp_path_retention_count",
  198. help="How many sessions should we keep the `tmp_path` directories, according to `tmp_path_retention_policy`.",
  199. default="3",
  200. # NOTE: Would have been better as an `int` but can't change it now.
  201. type="string",
  202. )
  203. parser.addini(
  204. "tmp_path_retention_policy",
  205. help="Controls which directories created by the `tmp_path` fixture are kept around, based on test outcome. "
  206. "(all/failed/none)",
  207. type="string",
  208. default="all",
  209. )
  210. @fixture(scope="session")
  211. def tmp_path_factory(request: FixtureRequest) -> TempPathFactory:
  212. """Return a :class:`pytest.TempPathFactory` instance for the test session."""
  213. # Set dynamically by pytest_configure() above.
  214. return request.config._tmp_path_factory # type: ignore
  215. def _mk_tmp(request: FixtureRequest, factory: TempPathFactory) -> Path:
  216. name = request.node.name
  217. name = re.sub(r"[\W]", "_", name)
  218. MAXVAL = 30
  219. name = name[:MAXVAL]
  220. return factory.mktemp(name, numbered=True)
  221. @fixture
  222. def tmp_path(
  223. request: FixtureRequest, tmp_path_factory: TempPathFactory
  224. ) -> Generator[Path]:
  225. """Return a temporary directory (as :class:`pathlib.Path` object)
  226. which is unique to each test function invocation.
  227. The temporary directory is created as a subdirectory
  228. of the base temporary directory, with configurable retention,
  229. as discussed in :ref:`temporary directory location and retention`.
  230. """
  231. path = _mk_tmp(request, tmp_path_factory)
  232. yield path
  233. # Remove the tmpdir if the policy is "failed" and the test passed.
  234. policy = tmp_path_factory._retention_policy
  235. result_dict = request.node.stash[tmppath_result_key]
  236. if policy == "failed" and result_dict.get("call", True):
  237. # We do a "best effort" to remove files, but it might not be possible due to some leaked resource,
  238. # permissions, etc, in which case we ignore it.
  239. rmtree(path, ignore_errors=True)
  240. del request.node.stash[tmppath_result_key]
  241. def pytest_sessionfinish(session, exitstatus: int | ExitCode):
  242. """After each session, remove base directory if all the tests passed,
  243. the policy is "failed", and the basetemp is not specified by a user.
  244. """
  245. tmp_path_factory: TempPathFactory = session.config._tmp_path_factory
  246. basetemp = tmp_path_factory._basetemp
  247. if basetemp is None:
  248. return
  249. policy = tmp_path_factory._retention_policy
  250. if (
  251. exitstatus == 0
  252. and policy == "failed"
  253. and tmp_path_factory._given_basetemp is None
  254. ):
  255. if basetemp.is_dir():
  256. # We do a "best effort" to remove files, but it might not be possible due to some leaked resource,
  257. # permissions, etc, in which case we ignore it.
  258. rmtree(basetemp, ignore_errors=True)
  259. # Remove dead symlinks.
  260. if basetemp.is_dir():
  261. cleanup_dead_symlinks(basetemp)
  262. @hookimpl(wrapper=True, tryfirst=True)
  263. def pytest_runtest_makereport(
  264. item: Item, call
  265. ) -> Generator[None, TestReport, TestReport]:
  266. rep = yield
  267. assert rep.when is not None
  268. empty: dict[str, bool] = {}
  269. item.stash.setdefault(tmppath_result_key, empty)[rep.when] = rep.passed
  270. return rep