Store grammar pickle caches in CACHE_DIR

Fixes pytest-dev#192

Fixes pytest-dev#203

Author: ambv

black.py

@@ -46,9 +46,10 @@
 
 __version__ = "18.5b0"
 DEFAULT_LINE_LENGTH = 88
+CACHE_DIR = Path(user_cache_dir("black", version=__version__))
+
 
 # types
-syms = pygram.python_symbols
 FileContent = str
 Encoding = str
 Depth = int
@@ -65,6 +66,9 @@
 out = partial(click.secho, bold=True, err=True)
 err = partial(click.secho, fg="red", err=True)
 
+pygram.initialize(CACHE_DIR)
+syms = pygram.python_symbols
+
 
 class NothingChanged(UserWarning):
     """Raised by :func:`format_file` when reformatted code is the same as source."""
@@ -3051,9 +3055,6 @@ def can_omit_invisible_parens(line: Line, line_length: int) -> bool:
     return False
 
 
-CACHE_DIR = Path(user_cache_dir("black", version=__version__))
-
-
 def get_cache_file(line_length: int) -> Path:
     return CACHE_DIR / f"cache.{line_length}.pickle"
 

blib2to3/pgen2/driver.py

@@ -149,11 +149,15 @@ def _partially_consume_prefix(self, prefix, column):
         return ''.join(lines), current_line
 
 
-def _generate_pickle_name(gt):
+def _generate_pickle_name(gt, cache_dir=None):
     head, tail = os.path.splitext(gt)
     if tail == ".txt":
         tail = ""
-    return head + tail + ".".join(map(str, sys.version_info)) + ".pickle"
+    name = head + tail + ".".join(map(str, sys.version_info)) + ".pickle"
+    if cache_dir:
+        return os.path.join(cache_dir, os.path.basename(name))
+    else:
+        return name
 
 
 def load_grammar(gt="Grammar.txt", gp=None,
@@ -186,7 +190,7 @@ def _newer(a, b):
     return os.path.getmtime(a) >= os.path.getmtime(b)
 
 
-def load_packaged_grammar(package, grammar_source):
+def load_packaged_grammar(package, grammar_source, cache_dir=None):
     """Normally, loads a pickled grammar by doing
         pkgutil.get_data(package, pickled_grammar)
     where *pickled_grammar* is computed from *grammar_source* by adding the
@@ -198,8 +202,9 @@ def load_packaged_grammar(package, grammar_source):
 
     """
     if os.path.isfile(grammar_source):
-        return load_grammar(grammar_source)
-    pickled_name = _generate_pickle_name(os.path.basename(grammar_source))
+        gp = _generate_pickle_name(grammar_source, cache_dir) if cache_dir else None
+        return load_grammar(grammar_source, gp=gp)
+    pickled_name = _generate_pickle_name(os.path.basename(grammar_source), cache_dir)
     data = pkgutil.get_data(package, pickled_name)
     g = grammar.Grammar()
     g.loads(data)

blib2to3/pygram.py

@@ -29,19 +29,29 @@ def __init__(self, grammar):
             setattr(self, name, symbol)
 
 
-# Python 2
-python_grammar = driver.load_packaged_grammar("blib2to3", _GRAMMAR_FILE)
-
-python_symbols = Symbols(python_grammar)
-
-# Python 2 + from __future__ import print_function
-python_grammar_no_print_statement = python_grammar.copy()
-del python_grammar_no_print_statement.keywords["print"]
-
-# Python 3
-python_grammar_no_print_statement_no_exec_statement = python_grammar.copy()
-del python_grammar_no_print_statement_no_exec_statement.keywords["print"]
-del python_grammar_no_print_statement_no_exec_statement.keywords["exec"]
-
-pattern_grammar = driver.load_packaged_grammar("blib2to3", _PATTERN_GRAMMAR_FILE)
-pattern_symbols = Symbols(pattern_grammar)
+def initialize(cache_dir=None):
+    global python_grammar
+    global python_grammar_no_print_statement
+    global python_grammar_no_print_statement_no_exec_statement
+    global python_symbols
+    global pattern_grammar
+    global pattern_symbols
+
+    # Python 2
+    python_grammar = driver.load_packaged_grammar("blib2to3", _GRAMMAR_FILE,
+                                                  cache_dir)
+
+    python_symbols = Symbols(python_grammar)
+
+    # Python 2 + from __future__ import print_function
+    python_grammar_no_print_statement = python_grammar.copy()
+    del python_grammar_no_print_statement.keywords["print"]
+
+    # Python 3
+    python_grammar_no_print_statement_no_exec_statement = python_grammar.copy()
+    del python_grammar_no_print_statement_no_exec_statement.keywords["print"]
+    del python_grammar_no_print_statement_no_exec_statement.keywords["exec"]
+
+    pattern_grammar = driver.load_packaged_grammar("blib2to3", _PATTERN_GRAMMAR_FILE,
+                                                   cache_dir)
+    pattern_symbols = Symbols(pattern_grammar)

blib2to3/pygram.pyi

@@ -1,6 +1,7 @@
 # Stubs for lib2to3.pygram (Python 3.6)
 
-from typing import Any
+import os
+from typing import Any, Union
 from blib2to3.pgen2.grammar import Grammar
 
 class Symbols:
@@ -119,3 +120,5 @@ python_grammar_no_print_statement: Grammar
 python_grammar_no_print_statement_no_exec_statement: Grammar
 python_grammar_no_exec_statement: Grammar
 pattern_grammar: Grammar
+
+def initialize(cache_dir: Union[str, os.PathLike, None]) -> None: ...