bpo-24960: use pkgutil.get_data in lib2to3 to read pickled grammar files (#4977)

This is more complicated than it should be because we need to preserve the
useful mtime-based regeneration feature that lib2to3.pgen2.driver.load_grammar
has. We only look for the pickled grammar file with pkgutil.get_data and only if
the source file does not exist.

4 files changed, 42 insertions, 2 deletions

diff --git a/Lib/lib2to3/pgen2/driver.py b/Lib/lib2to3/pgen2/driver.py
index e5e4824..95c6692 100644
--- a/Lib/lib2to3/pgen2/driver.py
+++ b/Lib/lib2to3/pgen2/driver.py
@@ -20,6 +20,7 @@ import codecs
 import io
 import os
 import logging
+import pkgutil
 import sys
 
 # Pgen imports
@@ -140,6 +141,26 @@ def _newer(a, b):
     return os.path.getmtime(a) >= os.path.getmtime(b)
 
 
+def load_packaged_grammar(package, grammar_source):
+    """Normally, loads a pickled grammar by doing
+        pkgutil.get_data(package, pickled_grammar)
+    where *pickled_grammar* is computed from *grammar_source* by adding the
+    Python version and using a ``.pickle`` extension.
+
+    However, if *grammar_source* is an extant file, load_grammar(grammar_source)
+    is called instead. This facilities using a packaged grammar file when needed
+    but preserves load_grammar's automatic regeneration behavior when possible.
+
+    """
+    if os.path.isfile(grammar_source):
+        return load_grammar(grammar_source)
+    pickled_name = _generate_pickle_name(os.path.basename(grammar_source))
+    data = pkgutil.get_data(package, pickled_name)
+    g = grammar.Grammar()
+    g.loads(data)
+    return g
+
+
 def main(*args):
     """Main program, when run as a script: produce grammar pickle files.
 
diff --git a/Lib/lib2to3/pgen2/grammar.py b/Lib/lib2to3/pgen2/grammar.py
index c10dcfa..088c58b 100644
--- a/Lib/lib2to3/pgen2/grammar.py
+++ b/Lib/lib2to3/pgen2/grammar.py
@@ -108,6 +108,10 @@ class Grammar(object):
             d = pickle.load(f)
         self.__dict__.update(d)
 
+    def loads(self, pkl):
+        """Load the grammar tables from a pickle bytes object."""
+        self.__dict__.update(pickle.loads(pkl))
+
     def copy(self):
         """
         Copy the grammar.
diff --git a/Lib/lib2to3/pygram.py b/Lib/lib2to3/pygram.py
index 01fa108..919624e 100644
--- a/Lib/lib2to3/pygram.py
+++ b/Lib/lib2to3/pygram.py
@@ -29,12 +29,12 @@ class Symbols(object):
             setattr(self, name, symbol)
 
 
-python_grammar = driver.load_grammar(_GRAMMAR_FILE)
+python_grammar = driver.load_packaged_grammar("lib2to3", _GRAMMAR_FILE)
 
 python_symbols = Symbols(python_grammar)
 
 python_grammar_no_print_statement = python_grammar.copy()
 del python_grammar_no_print_statement.keywords["print"]
 
-pattern_grammar = driver.load_grammar(_PATTERN_GRAMMAR_FILE)
+pattern_grammar = driver.load_packaged_grammar("lib2to3", _PATTERN_GRAMMAR_FILE)
 pattern_symbols = Symbols(pattern_grammar)
diff --git a/Lib/lib2to3/tests/test_parser.py b/Lib/lib2to3/tests/test_parser.py
index dc94a69..ed1ef83 100644
--- a/Lib/lib2to3/tests/test_parser.py
+++ b/Lib/lib2to3/tests/test_parser.py
@@ -12,7 +12,10 @@ from .support import driver
 from test.support import verbose
 
 # Python imports
+import importlib
+import operator
 import os
+import pickle
 import shutil
 import subprocess
 import sys
@@ -99,6 +102,18 @@ pgen2_driver.load_grammar(%r, save=True, force=True)
         finally:
             shutil.rmtree(tmpdir)
 
+    def test_load_packaged_grammar(self):
+        modname = __name__ + '.load_test'
+        class MyLoader:
+            def get_data(self, where):
+                return pickle.dumps({'elephant': 19})
+        class MyModule:
+            __file__ = 'parsertestmodule'
+            __spec__ = importlib.util.spec_from_loader(modname, MyLoader())
+        sys.modules[modname] = MyModule()
+        self.addCleanup(operator.delitem, sys.modules, modname)
+        g = pgen2_driver.load_packaged_grammar(modname, 'Grammar.txt')
+        self.assertEqual(g.elephant, 19)
 
 
 class GrammarTest(support.TestCase):