[Python-checkins] [2.7] bpo-24960: use pkgutil.get_data in lib2to3 to read pickled grammar files (GH-4977) (#4980)

[Benjamin Peterson]

https://github.com/python/cpython/commit/770a802df7057bcf0d72bbc01026fac8d396e1b3
commit: 770a802df7057bcf0d72bbc01026fac8d396e1b3
branch: 2.7
author: Benjamin Peterson <benjamin at python.org>
committer: GitHub <noreply at github.com>
date: 2017-12-22T12:52:10-08:00
summary:

[2.7] bpo-24960: use pkgutil.get_data in lib2to3 to read pickled grammar files (GH-4977) (#4980)

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..
(cherry picked from commit 8a5877165e993afb2633cd48da5222326d3f6e0e)

files:
A Misc/NEWS.d/next/Tools-Demos/2017-12-22-09-25-51.bpo-24960.TGdAgO.rst
M Lib/lib2to3/pgen2/driver.py
M Lib/lib2to3/pgen2/grammar.py
M Lib/lib2to3/pygram.py
M Lib/lib2to3/tests/test_parser.py

diff --git a/Lib/lib2to3/pgen2/driver.py b/Lib/lib2to3/pgen2/driver.py
index ce601bb04f9..e8d10e47dba 100644
--- a/Lib/lib2to3/pgen2/driver.py
+++ b/Lib/lib2to3/pgen2/driver.py
@@ -19,6 +19,7 @@
 import codecs
 import os
 import logging
+import pkgutil
 import StringIO
 import sys
 
@@ -143,6 +144,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 75255e9c013..0b6d86b6799 100644
--- a/Lib/lib2to3/pgen2/grammar.py
+++ b/Lib/lib2to3/pgen2/grammar.py
@@ -109,6 +109,10 @@ def load(self, filename):
         f.close()
         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 621ff24c954..7e67e4a8678 100644
--- a/Lib/lib2to3/pygram.py
+++ b/Lib/lib2to3/pygram.py
@@ -29,12 +29,12 @@ def __init__(self, grammar):
             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 ebf84418fe6..d2254f13e09 100644
--- a/Lib/lib2to3/tests/test_parser.py
+++ b/Lib/lib2to3/tests/test_parser.py
@@ -11,11 +11,14 @@
 from .support import driver, test_dir
 
 # Python imports
+import operator
 import os
+import pickle
 import shutil
 import subprocess
 import sys
 import tempfile
+import types
 import unittest
 
 # Local imports
@@ -97,6 +100,18 @@ def test_load_grammar_from_subprocess(self):
         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(types.ModuleType):
+            __file__ = 'parsertestmodule'
+            __loader__ = MyLoader()
+        sys.modules[modname] = MyModule(modname)
+        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):
diff --git a/Misc/NEWS.d/next/Tools-Demos/2017-12-22-09-25-51.bpo-24960.TGdAgO.rst b/Misc/NEWS.d/next/Tools-Demos/2017-12-22-09-25-51.bpo-24960.TGdAgO.rst
new file mode 100644
index 00000000000..6f9d83c5772
--- /dev/null
+++ b/Misc/NEWS.d/next/Tools-Demos/2017-12-22-09-25-51.bpo-24960.TGdAgO.rst
@@ -0,0 +1,3 @@
+2to3 and lib2to3 can now read pickled grammar files using pkgutil.get_data()
+rather than probing the filesystem. This lets 2to3 and lib2to3 work when run
+from a zipfile.