Support sys.pycache_prefix on py38

Fix #4730

Author: nicoddemus

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 sys
 import tokenize
 import types
+from pathlib import Path
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -30,7 +31,7 @@
 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 +104,7 @@ def create_module(self, spec):
         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 +118,15 @@ def exec_module(self, module):
         # 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)
             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 +140,7 @@ def exec_module(self, module):
                 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 +259,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(str(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,14 +270,15 @@ 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
 
 
 def _rewrite_test(fn, config):
     """read and rewrite *fn* and return the code object."""
+    fn = str(fn)
     stat = os.stat(fn)
     with open(fn, "rb") as f:
         source = f.read()
@@ -292,12 +294,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(str(pyc), "rb")
     except IOError:
         return None
     with fp:
         try:
-            stat_result = os.stat(source)
+            stat_result = os.stat(str(source))
             mtime = int(stat_result.st_mtime)
             size = stat_result.st_size
             data = fp.read(12)
@@ -749,7 +751,7 @@ def visit_Assert(self, assert_):
                     "assertion is always true, perhaps remove parentheses?"
                 ),
                 category=None,
-                filename=self.module_path,
+                filename=str(self.module_path),
                 lineno=assert_.lineno,
             )
 
@@ -872,7 +874,7 @@ def warn_about_none_ast(self, node, module_path, lineno):
     lineno={lineno},
 )
             """.format(
-                filename=module_path, lineno=lineno
+                filename=str(module_path), lineno=lineno
             )
         ).body
         return ast.If(val_is_none, send_warning, [])
@@ -1021,9 +1023,9 @@ def visit_Compare(self, comp: ast.Compare):
 def try_mkdir(cache_dir):
     """Attempts to create the given directory, returns True if successful"""
     try:
-        os.mkdir(cache_dir)
+        os.makedirs(str(cache_dir))
     except FileExistsError:
-        # Either the __pycache__ directory already exists (the
+        # 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
@@ -1039,3 +1041,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__"

testing/test_assertrewrite.py

@@ -9,6 +9,7 @@
 import textwrap
 import zipfile
 from functools import partial
+from pathlib import Path
 
 import py
 
@@ -17,6 +18,8 @@
 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
@@ -1564,7 +1567,7 @@ def test_try_mkdir(monkeypatch, tmp_path):
     assert try_mkdir(str(p))
 
     # monkeypatch to simulate all error situations
-    def fake_mkdir(p, *, exc):
+    def fake_mkdir(p, mode, *, exc):
         assert isinstance(p, str)
         raise exc
 
@@ -1589,3 +1592,59 @@ def fake_mkdir(p, *, exc):
     with pytest.raises(OSError) as exc_info:
         try_mkdir(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()