bpo-33499: Add PYTHONPYCACHEPREFIX env var for alt bytecode cache location. (GH-6834)

In some development setups it is inconvenient or impossible to write bytecode
caches to the code tree, but the bytecode caches are still useful. The
PYTHONPYCACHEPREFIX environment variable allows specifying an alternate
location for cached bytecode files, within which a directory tree mirroring the code
tree will be created. This cache tree is then used (for both reading and writing)
instead of the local `__pycache__` subdirectory within each source directory.

Exposed at runtime as sys.pycache_prefix (defaulting to None), and can
be set from the CLI as "-X pycache_prefix=path".

Patch by Carl Meyer.

2 files changed, 95 insertions, 11 deletions

diff --git a/Lib/test/test_importlib/test_util.py b/Lib/test/test_importlib/test_util.py
index 0946b34..b319e18 100644
--- a/Lib/test/test_importlib/test_util.py
+++ b/Lib/test/test_importlib/test_util.py
@@ -1,9 +1,10 @@
-from . import util
+from . import util
 abc = util.import_importlib('importlib.abc')
 init = util.import_importlib('importlib')
 machinery = util.import_importlib('importlib.machinery')
 importlib_util = util.import_importlib('importlib.util')
 
+import contextlib
 import importlib.util
 import os
 import pathlib
@@ -12,6 +13,7 @@ import sys
 from test import support
 import types
 import unittest
+import unittest.mock
 import warnings
 
 
@@ -557,8 +559,8 @@ class PEP3147Tests:
 
     tag = sys.implementation.cache_tag
 
-    @unittest.skipUnless(sys.implementation.cache_tag is not None,
-                         'requires sys.implementation.cache_tag not be None')
+    @unittest.skipIf(sys.implementation.cache_tag is None,
+                     'requires sys.implementation.cache_tag not be None')
     def test_cache_from_source(self):
         # Given the path to a .py file, return the path to its PEP 3147
         # defined .pyc file (i.e. under __pycache__).
@@ -678,8 +680,8 @@ class PEP3147Tests:
             self.util.cache_from_source('\\foo\\bar\\baz/qux.py', optimization=''),
             '\\foo\\bar\\baz\\__pycache__\\qux.{}.pyc'.format(self.tag))
 
-    @unittest.skipUnless(sys.implementation.cache_tag is not None,
-                         'requires sys.implementation.cache_tag not be None')
+    @unittest.skipIf(sys.implementation.cache_tag is None,
+                     'requires sys.implementation.cache_tag not be None')
     def test_source_from_cache_path_like_arg(self):
         path = pathlib.PurePath('foo', 'bar', 'baz', 'qux.py')
         expect = os.path.join('foo', 'bar', 'baz', '__pycache__',
@@ -687,9 +689,8 @@ class PEP3147Tests:
         self.assertEqual(self.util.cache_from_source(path, optimization=''),
                          expect)
 
-    @unittest.skipUnless(sys.implementation.cache_tag is not None,
-                         'requires sys.implementation.cache_tag to not be '
-                         'None')
+    @unittest.skipIf(sys.implementation.cache_tag is None,
+                     'requires sys.implementation.cache_tag to not be None')
     def test_source_from_cache(self):
         # Given the path to a PEP 3147 defined .pyc file, return the path to
         # its source.  This tests the good path.
@@ -749,15 +750,87 @@ class PEP3147Tests:
         with self.assertRaises(ValueError):
             self.util.source_from_cache(path)
 
-    @unittest.skipUnless(sys.implementation.cache_tag is not None,
-                         'requires sys.implementation.cache_tag to not be '
-                         'None')
+    @unittest.skipIf(sys.implementation.cache_tag is None,
+                     'requires sys.implementation.cache_tag to not be None')
     def test_source_from_cache_path_like_arg(self):
         path = pathlib.PurePath('foo', 'bar', 'baz', '__pycache__',
                                 'qux.{}.pyc'.format(self.tag))
         expect = os.path.join('foo', 'bar', 'baz', 'qux.py')
         self.assertEqual(self.util.source_from_cache(path), expect)
 
+    @unittest.skipIf(sys.implementation.cache_tag is None,
+                     'requires sys.implementation.cache_tag to not be None')
+    def test_cache_from_source_respects_pycache_prefix(self):
+        # If pycache_prefix is set, cache_from_source will return a bytecode
+        # path inside that directory (in a subdirectory mirroring the .py file's
+        # path) rather than in a __pycache__ dir next to the py file.
+        pycache_prefixes = [
+            os.path.join(os.path.sep, 'tmp', 'bytecode'),
+            os.path.join(os.path.sep, 'tmp', '\u2603'),  # non-ASCII in path!
+            os.path.join(os.path.sep, 'tmp', 'trailing-slash') + os.path.sep,
+        ]
+        drive = ''
+        if os.name == 'nt':
+            drive = 'C:'
+            pycache_prefixes = [
+                f'{drive}{prefix}' for prefix in pycache_prefixes]
+            pycache_prefixes += [r'\\?\C:\foo', r'\\localhost\c$\bar']
+        for pycache_prefix in pycache_prefixes:
+            with self.subTest(path=pycache_prefix):
+                path = drive + os.path.join(
+                    os.path.sep, 'foo', 'bar', 'baz', 'qux.py')
+                expect = os.path.join(
+                    pycache_prefix, 'foo', 'bar', 'baz',
+                    'qux.{}.pyc'.format(self.tag))
+                with util.temporary_pycache_prefix(pycache_prefix):
+                    self.assertEqual(
+                        self.util.cache_from_source(path, optimization=''),
+                        expect)
+
+    @unittest.skipIf(sys.implementation.cache_tag is None,
+                     'requires sys.implementation.cache_tag to not be None')
+    def test_cache_from_source_respects_pycache_prefix_relative(self):
+        # If the .py path we are given is relative, we will resolve to an
+        # absolute path before prefixing with pycache_prefix, to avoid any
+        # possible ambiguity.
+        pycache_prefix = os.path.join(os.path.sep, 'tmp', 'bytecode')
+        path = os.path.join('foo', 'bar', 'baz', 'qux.py')
+        root = os.path.splitdrive(os.getcwd())[0] + os.path.sep
+        expect = os.path.join(
+            pycache_prefix,
+            os.path.relpath(os.getcwd(), root),
+            'foo', 'bar', 'baz', f'qux.{self.tag}.pyc')
+        with util.temporary_pycache_prefix(pycache_prefix):
+            self.assertEqual(
+                self.util.cache_from_source(path, optimization=''),
+                expect)
+
+    @unittest.skipIf(sys.implementation.cache_tag is None,
+                     'requires sys.implementation.cache_tag to not be None')
+    def test_source_from_cache_inside_pycache_prefix(self):
+        # If pycache_prefix is set and the cache path we get is inside it,
+        # we return an absolute path to the py file based on the remainder of
+        # the path within pycache_prefix.
+        pycache_prefix = os.path.join(os.path.sep, 'tmp', 'bytecode')
+        path = os.path.join(pycache_prefix, 'foo', 'bar', 'baz',
+                            f'qux.{self.tag}.pyc')
+        expect = os.path.join(os.path.sep, 'foo', 'bar', 'baz', 'qux.py')
+        with util.temporary_pycache_prefix(pycache_prefix):
+            self.assertEqual(self.util.source_from_cache(path), expect)
+
+    @unittest.skipIf(sys.implementation.cache_tag is None,
+                     'requires sys.implementation.cache_tag to not be None')
+    def test_source_from_cache_outside_pycache_prefix(self):
+        # If pycache_prefix is set but the cache path we get is not inside
+        # it, just ignore it and handle the cache path according to the default
+        # behavior.
+        pycache_prefix = os.path.join(os.path.sep, 'tmp', 'bytecode')
+        path = os.path.join('foo', 'bar', 'baz', '__pycache__',
+                            f'qux.{self.tag}.pyc')
+        expect = os.path.join('foo', 'bar', 'baz', 'qux.py')
+        with util.temporary_pycache_prefix(pycache_prefix):
+            self.assertEqual(self.util.source_from_cache(path), expect)
+
 
 (Frozen_PEP3147Tests,
  Source_PEP3147Tests
diff --git a/Lib/test/test_importlib/util.py b/Lib/test/test_importlib/util.py
index b0badeb..196ea1c 100644
--- a/Lib/test/test_importlib/util.py
+++ b/Lib/test/test_importlib/util.py
@@ -320,6 +320,17 @@ def ensure_bytecode_path(bytecode_path):
 
 
 @contextlib.contextmanager
+def temporary_pycache_prefix(prefix):
+    """Adjust and restore sys.pycache_prefix."""
+    _orig_prefix = sys.pycache_prefix
+    sys.pycache_prefix = prefix
+    try:
+        yield
+    finally:
+        sys.pycache_prefix = _orig_prefix
+
+
+@contextlib.contextmanager
 def create_modules(*names):
     """Temporarily create each named module with an attribute (named 'attr')
     that contains the name passed into the context manager that caused the