Allow setting custom cache directory on all platforms with environmen…

CHANGES.md

@@ -20,6 +20,8 @@
 - Tuple unpacking on `return` and `yield` constructs now implies 3.8+ (#2700)
 - Unparenthesized tuples on annotated assignments (e.g
   `values: Tuple[int, ...] = 1, 2, 3`) now implies 3.8+ (#2708)
+- Allow setting custom cache directory on all platforms with environment variable
+  `BLACK_CACHE_DIR` (#2739).
 - Text coloring added in the final statistics (#2712)
 - For stubs, one blank line between class attributes and methods is now kept if there's
   at least one pre-existing blank line (#2736)

docs/contributing/reference/reference_functions.rst

@@ -96,6 +96,8 @@ Caching
 
 .. autofunction:: black.cache.filter_cached
 
+.. autofunction:: black.cache.get_cache_dir
+
 .. autofunction:: black.cache.get_cache_file
 
 .. autofunction:: black.cache.get_cache_info

docs/usage_and_configuration/file_collection_and_discovery.md

@@ -22,10 +22,12 @@ run. The file is non-portable. The standard location on common operating systems
 `file-mode` is an int flag that determines whether the file was formatted as 3.6+ only,
 as .pyi, and whether string normalization was omitted.
 
-To override the location of these files on macOS or Linux, set the environment variable
-`XDG_CACHE_HOME` to your preferred location. For example, if you want to put the cache
-in the directory you're running _Black_ from, set `XDG_CACHE_HOME=.cache`. _Black_ will
-then write the above files to `.cache/black/<version>/`.
+To override the location of these files on all systems, set the environment variable
+`BLACK_CACHE_DIR` to the preferred location. Alternatively on macOS and Linux, set
+`XDG_CACHE_HOME` to you your preferred location. For example, if you want to put the
+cache in the directory you're running _Black_ from, set `BLACK_CACHE_DIR=.cache/black`.
+_Black_ will then write the above files to `.cache/black`. Note that `BLACK_CACHE_DIR`
+will take precedence over `XDG_CACHE_HOME` if both are set.
 
 ## .gitignore
 

src/black/cache.py

@@ -20,7 +20,23 @@
 Cache = Dict[str, CacheInfo]
 
 
-CACHE_DIR = Path(user_cache_dir("black", version=__version__))
+def get_cache_dir() -> Path:
+    """Get the cache directory used by black.
+
+    Users can customize this directory on all systems using `BLACK_CACHE_DIR`
+    environment variable. By default, the cache directory is the user cache directory
+    under the black application.
+
+    This result is immediately set to a constant `black.cache.CACHE_DIR` as to avoid
+    repeated calls.
+    """
+    # NOTE: Function mostly exists as a clean way to test getting the cache directory.
+    default_cache_dir = user_cache_dir("black", version=__version__)
+    cache_dir = Path(os.environ.get("BLACK_CACHE_DIR", default_cache_dir))
+    return cache_dir
+
+
+CACHE_DIR = get_cache_dir()
 
 
 def read_cache(mode: Mode) -> Cache:

tests/test_black.py

@@ -40,7 +40,7 @@
 import black.files
 from black import Feature, TargetVersion
 from black import re_compile_maybe_verbose as compile_pattern
-from black.cache import get_cache_file
+from black.cache import get_cache_dir, get_cache_file
 from black.debug import DebugVisitor
 from black.output import color_diff, diff
 from black.report import Report
@@ -1601,6 +1601,33 @@ def test_equivalency_ast_parse_failure_includes_error(self) -> None:
 
 
 class TestCaching:
+    def test_get_cache_dir(
+        self,
+        tmp_path: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        # Create multiple cache directories
+        workspace1 = tmp_path / "ws1"
+        workspace1.mkdir()
+        workspace2 = tmp_path / "ws2"
+        workspace2.mkdir()
+
+        # Force user_cache_dir to use the temporary directory for easier assertions
+        patch_user_cache_dir = patch(
+            target="black.cache.user_cache_dir",
+            autospec=True,
+            return_value=str(workspace1),
+        )
+
+        # If BLACK_CACHE_DIR is not set, use user_cache_dir
+        monkeypatch.delenv("BLACK_CACHE_DIR", raising=False)
+        with patch_user_cache_dir:
+            assert get_cache_dir() == workspace1
+
+        # If it is set, use the path provided in the env var.
+        monkeypatch.setenv("BLACK_CACHE_DIR", str(workspace2))
+        assert get_cache_dir() == workspace2
+
     def test_cache_broken_file(self) -> None:
         mode = DEFAULT_MODE
         with cache_dir() as workspace: