Support sys.pycache_prefix on py38 (#5864)

Support sys.pycache_prefix on py38

changelog/4730.feature.rst

@@ -0,0 +1,3 @@
+When ``sys.pycache_prefix`` (Python 3.8+) is set, it will be used by pytest to cache test files changed by the assertion rewriting mechanism.
+
+This makes it easier to benefit of cached ``.pyc`` files even on file systems without permissions.

src/_pytest/assertion/rewrite.py

@@ -13,6 +13,7 @@ import struct
 import sys
 import tokenize
 import types
+from pathlib import Path
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -27,10 +28,11 @@ from _pytest.assertion import util
 from _pytest.assertion.util import (  # noqa: F401
     format_explanation as _format_explanation,
 )
+from _pytest.compat import fspath
 from _pytest.pathlib import fnmatch_ex
 from _pytest.pathlib import PurePath
 
-# pytest caches rewritten pycs in __pycache__.
+# pytest caches rewritten pycs in pycache dirs
 PYTEST_TAG = "{}-pytest-{}".format(sys.implementation.cache_tag, version)
 PYC_EXT = ".py" + (__debug__ and "c" or "o")
 PYC_TAIL = "." + PYTEST_TAG + PYC_EXT
@@ -103,7 +105,7 @@ class AssertionRewritingHook(importlib.abc.MetaPathFinder):
         return None  # default behaviour is fine
 
     def exec_module(self, module):
-        fn = module.__spec__.origin
+        fn = Path(module.__spec__.origin)
         state = self.config._assertstate
 
         self._rewritten_names.add(module.__name__)
@@ -117,15 +119,15 @@ class AssertionRewritingHook(importlib.abc.MetaPathFinder):
         # cached pyc is always a complete, valid pyc. Operations on it must be
         # atomic. POSIX's atomic rename comes in handy.
         write = not sys.dont_write_bytecode
-        cache_dir = os.path.join(os.path.dirname(fn), "__pycache__")
+        cache_dir = get_cache_dir(fn)
         if write:
-            ok = try_mkdir(cache_dir)
+            ok = try_makedirs(cache_dir)
             if not ok:
                 write = False
-                state.trace("read only directory: {}".format(os.path.dirname(fn)))
+                state.trace("read only directory: {}".format(cache_dir))
 
-        cache_name = os.path.basename(fn)[:-3] + PYC_TAIL
-        pyc = os.path.join(cache_dir, cache_name)
+        cache_name = fn.name[:-3] + PYC_TAIL
+        pyc = cache_dir / cache_name
         # Notice that even if we're in a read-only directory, I'm going
         # to check for a cached pyc. This may not be optimal...
         co = _read_pyc(fn, pyc, state.trace)
@@ -139,7 +141,7 @@ class AssertionRewritingHook(importlib.abc.MetaPathFinder):
                 finally:
                     self._writing_pyc = False
         else:
-            state.trace("found cached rewritten pyc for {!r}".format(fn))
+            state.trace("found cached rewritten pyc for {}".format(fn))
         exec(co, module.__dict__)
 
     def _early_rewrite_bailout(self, name, state):
@@ -258,7 +260,7 @@ def _write_pyc(state, co, source_stat, pyc):
     # (C)Python, since these "pycs" should never be seen by builtin
     # import. However, there's little reason deviate.
     try:
-        with atomicwrites.atomic_write(pyc, mode="wb", overwrite=True) as fp:
+        with atomicwrites.atomic_write(fspath(pyc), mode="wb", overwrite=True) as fp:
             fp.write(importlib.util.MAGIC_NUMBER)
             # as of now, bytecode header expects 32-bit numbers for size and mtime (#4903)
             mtime = int(source_stat.st_mtime) & 0xFFFFFFFF
@@ -269,7 +271,7 @@ def _write_pyc(state, co, source_stat, pyc):
     except EnvironmentError as e:
         state.trace("error writing pyc file at {}: errno={}".format(pyc, e.errno))
         # we ignore any failure to write the cache file
-        # there are many reasons, permission-denied, __pycache__ being a
+        # there are many reasons, permission-denied, pycache dir being a
         # file etc.
         return False
     return True
@@ -277,6 +279,7 @@ def _write_pyc(state, co, source_stat, pyc):
 
 def _rewrite_test(fn, config):
     """read and rewrite *fn* and return the code object."""
+    fn = fspath(fn)
     stat = os.stat(fn)
     with open(fn, "rb") as f:
         source = f.read()
@@ -292,12 +295,12 @@ def _read_pyc(source, pyc, trace=lambda x: None):
     Return rewritten code if successful or None if not.
     """
     try:
-        fp = open(pyc, "rb")
+        fp = open(fspath(pyc), "rb")
     except IOError:
         return None
     with fp:
         try:
-            stat_result = os.stat(source)
+            stat_result = os.stat(fspath(source))
             mtime = int(stat_result.st_mtime)
             size = stat_result.st_size
             data = fp.read(12)
@@ -749,7 +752,7 @@ class AssertionRewriter(ast.NodeVisitor):
                     "assertion is always true, perhaps remove parentheses?"
                 ),
                 category=None,
-                filename=self.module_path,
+                filename=fspath(self.module_path),
                 lineno=assert_.lineno,
             )
 
@@ -872,7 +875,7 @@ warn_explicit(
     lineno={lineno},
 )
             """.format(
-                filename=module_path, lineno=lineno
+                filename=fspath(module_path), lineno=lineno
             )
         ).body
         return ast.If(val_is_none, send_warning, [])
@@ -1018,18 +1021,15 @@ warn_explicit(
         return res, self.explanation_param(self.pop_format_context(expl_call))
 
 
-def try_mkdir(cache_dir):
-    """Attempts to create the given directory, returns True if successful"""
+def try_makedirs(cache_dir) -> bool:
+    """Attempts to create the given directory and sub-directories exist, returns True if
+    successful or it already exists"""
     try:
-        os.mkdir(cache_dir)
-    except FileExistsError:
-        # Either the __pycache__ directory already exists (the
-        # common case) or it's blocked by a non-dir node. In the
-        # latter case, we'll ignore it in _write_pyc.
-        return True
-    except (FileNotFoundError, NotADirectoryError):
-        # One of the path components was not a directory, likely
-        # because we're in a zip file.
+        os.makedirs(fspath(cache_dir), exist_ok=True)
+    except (FileNotFoundError, NotADirectoryError, FileExistsError):
+        # One of the path components was not a directory:
+        # - we're in a zip file
+        # - it is a file
         return False
     except PermissionError:
         return False
@@ -1039,3 +1039,17 @@ def try_mkdir(cache_dir):
             return False
         raise
     return True
+
+
+def get_cache_dir(file_path: Path) -> Path:
+    """Returns the cache directory to write .pyc files for the given .py file path"""
+    if sys.version_info >= (3, 8) and sys.pycache_prefix:
+        # given:
+        #   prefix = '/tmp/pycs'
+        #   path = '/home/user/proj/test_app.py'
+        # we want:
+        #   '/tmp/pycs/home/user/proj'
+        return Path(sys.pycache_prefix) / Path(*file_path.parts[1:-1])
+    else:
+        # classic pycache directory
+        return file_path.parent / "__pycache__"

src/_pytest/compat.py

@@ -4,6 +4,7 @@ python version compatibility code
 import functools
 import inspect
 import io
+import os
 import re
 import sys
 from contextlib import contextmanager
@@ -41,6 +42,19 @@ def _format_args(func):
 REGEX_TYPE = type(re.compile(""))
 
 
+if sys.version_info < (3, 6):
+
+    def fspath(p):
+        """os.fspath replacement, useful to point out when we should replace it by the
+        real function once we drop py35.
+        """
+        return str(p)
+
+
+else:
+    fspath = os.fspath
+
+
 def is_generator(func):
     genfunc = inspect.isgeneratorfunction(func)
     return genfunc and not iscoroutinefunction(func)

testing/test_assertrewrite.py

@@ -9,6 +9,7 @@ import sys
 import textwrap
 import zipfile
 from functools import partial
+from pathlib import Path
 
 import py
 
@@ -17,6 +18,8 @@ import pytest
 from _pytest.assertion import util
 from _pytest.assertion.rewrite import _get_assertion_exprs
 from _pytest.assertion.rewrite import AssertionRewritingHook
+from _pytest.assertion.rewrite import get_cache_dir
+from _pytest.assertion.rewrite import PYC_TAIL
 from _pytest.assertion.rewrite import PYTEST_TAG
 from _pytest.assertion.rewrite import rewrite_asserts
 from _pytest.main import ExitCode
@@ -1551,41 +1554,97 @@ def test_get_assertion_exprs(src, expected):
     assert _get_assertion_exprs(src) == expected
 
 
-def test_try_mkdir(monkeypatch, tmp_path):
-    from _pytest.assertion.rewrite import try_mkdir
+def test_try_makedirs(monkeypatch, tmp_path):
+    from _pytest.assertion.rewrite import try_makedirs
 
     p = tmp_path / "foo"
 
     # create
-    assert try_mkdir(str(p))
+    assert try_makedirs(str(p))
     assert p.is_dir()
 
     # already exist
-    assert try_mkdir(str(p))
+    assert try_makedirs(str(p))
 
     # monkeypatch to simulate all error situations
-    def fake_mkdir(p, *, exc):
+    def fake_mkdir(p, exist_ok=False, *, exc):
         assert isinstance(p, str)
         raise exc
 
-    monkeypatch.setattr(os, "mkdir", partial(fake_mkdir, exc=FileNotFoundError()))
-    assert not try_mkdir(str(p))
+    monkeypatch.setattr(os, "makedirs", partial(fake_mkdir, exc=FileNotFoundError()))
+    assert not try_makedirs(str(p))
 
-    monkeypatch.setattr(os, "mkdir", partial(fake_mkdir, exc=NotADirectoryError()))
-    assert not try_mkdir(str(p))
+    monkeypatch.setattr(os, "makedirs", partial(fake_mkdir, exc=NotADirectoryError()))
+    assert not try_makedirs(str(p))
 
-    monkeypatch.setattr(os, "mkdir", partial(fake_mkdir, exc=PermissionError()))
-    assert not try_mkdir(str(p))
+    monkeypatch.setattr(os, "makedirs", partial(fake_mkdir, exc=PermissionError()))
+    assert not try_makedirs(str(p))
 
     err = OSError()
     err.errno = errno.EROFS
-    monkeypatch.setattr(os, "mkdir", partial(fake_mkdir, exc=err))
-    assert not try_mkdir(str(p))
+    monkeypatch.setattr(os, "makedirs", partial(fake_mkdir, exc=err))
+    assert not try_makedirs(str(p))
 
     # unhandled OSError should raise
     err = OSError()
     err.errno = errno.ECHILD
-    monkeypatch.setattr(os, "mkdir", partial(fake_mkdir, exc=err))
+    monkeypatch.setattr(os, "makedirs", partial(fake_mkdir, exc=err))
     with pytest.raises(OSError) as exc_info:
-        try_mkdir(str(p))
+        try_makedirs(str(p))
     assert exc_info.value.errno == errno.ECHILD
+
+
+class TestPyCacheDir:
+    @pytest.mark.parametrize(
+        "prefix, source, expected",
+        [
+            ("c:/tmp/pycs", "d:/projects/src/foo.py", "c:/tmp/pycs/projects/src"),
+            (None, "d:/projects/src/foo.py", "d:/projects/src/__pycache__"),
+            ("/tmp/pycs", "/home/projects/src/foo.py", "/tmp/pycs/home/projects/src"),
+            (None, "/home/projects/src/foo.py", "/home/projects/src/__pycache__"),
+        ],
+    )
+    def test_get_cache_dir(self, monkeypatch, prefix, source, expected):
+        if prefix:
+            if sys.version_info < (3, 8):
+                pytest.skip("pycache_prefix not available in py<38")
+            monkeypatch.setattr(sys, "pycache_prefix", prefix)
+
+        assert get_cache_dir(Path(source)) == Path(expected)
+
+    @pytest.mark.skipif(
+        sys.version_info < (3, 8), reason="pycache_prefix not available in py<38"
+    )
+    def test_sys_pycache_prefix_integration(self, tmp_path, monkeypatch, testdir):
+        """Integration test for sys.pycache_prefix (#4730)."""
+        pycache_prefix = tmp_path / "my/pycs"
+        monkeypatch.setattr(sys, "pycache_prefix", str(pycache_prefix))
+        monkeypatch.setattr(sys, "dont_write_bytecode", False)
+
+        testdir.makepyfile(
+            **{
+                "src/test_foo.py": """
+                import bar
+                def test_foo():
+                    pass
+            """,
+                "src/bar/__init__.py": "",
+            }
+        )
+        result = testdir.runpytest()
+        assert result.ret == 0
+
+        test_foo = Path(testdir.tmpdir) / "src/test_foo.py"
+        bar_init = Path(testdir.tmpdir) / "src/bar/__init__.py"
+        assert test_foo.is_file()
+        assert bar_init.is_file()
+
+        # test file: rewritten, custom pytest cache tag
+        test_foo_pyc = get_cache_dir(test_foo) / ("test_foo" + PYC_TAIL)
+        assert test_foo_pyc.is_file()
+
+        # normal file: not touched by pytest, normal cache tag
+        bar_init_pyc = get_cache_dir(bar_init) / "__init__.{cache_tag}.pyc".format(
+            cache_tag=sys.implementation.cache_tag
+        )
+        assert bar_init_pyc.is_file()