bpo-45324: Capture data in FrozenImporter.find_spec() to use in exec_module(). (gh-28633)

Before this change we end up duplicating effort and throwing away data in FrozenImporter.find_spec().  Now we do the work once in find_spec() and the only thing we do in FrozenImporter.exec_module() is turn the raw frozen data into a code object and then exec it.

We've added _imp.find_frozen(), add an arg to _imp.get_frozen_object(), and updated FrozenImporter.  We've also moved some code around to reduce duplication, get a little more consistency in outcomes, and be more efficient.

Note that this change is mostly necessary if we want to set __file__ on frozen stdlib modules. (See https://bugs.python.org/issue21736.)

https://bugs.python.org/issue45324

2 files changed, 60 insertions, 23 deletions

diff --git a/Lib/test/test_importlib/frozen/test_finder.py b/Lib/test/test_importlib/frozen/test_finder.py
index 7d43fb0..23d1bf7 100644
--- a/Lib/test/test_importlib/frozen/test_finder.py
+++ b/Lib/test/test_importlib/frozen/test_finder.py
@@ -1,13 +1,15 @@
 from .. import abc
-import os.path
 from .. import util
 
 machinery = util.import_importlib('importlib.machinery')
 
+import _imp
+import marshal
+import os.path
 import unittest
 import warnings
 
-from test.support import import_helper
+from test.support import import_helper, REPO_ROOT
 
 
 class FindSpecTests(abc.FinderTests):
@@ -19,39 +21,67 @@ class FindSpecTests(abc.FinderTests):
         with import_helper.frozen_modules():
             return finder.find_spec(name, **kwargs)
 
-    def check(self, spec, name):
+    def check_basic(self, spec, name, ispkg=False):
         self.assertEqual(spec.name, name)
         self.assertIs(spec.loader, self.machinery.FrozenImporter)
         self.assertEqual(spec.origin, 'frozen')
         self.assertFalse(spec.has_location)
+        if ispkg:
+            self.assertIsNotNone(spec.submodule_search_locations)
+        else:
+            self.assertIsNone(spec.submodule_search_locations)
+        self.assertIsNotNone(spec.loader_state)
+
+    def check_search_location(self, spec, source=None):
+        # Frozen packages do not have any path entries.
+        # (See https://bugs.python.org/issue21736.)
+        expected = []
+        self.assertListEqual(spec.submodule_search_locations, expected)
+
+    def check_data(self, spec, source=None, ispkg=None):
+        with import_helper.frozen_modules():
+            expected = _imp.get_frozen_object(spec.name)
+        data = spec.loader_state
+        # We can't compare the marshaled data directly because
+        # marshal.dumps() would mark "expected" as a ref, which slightly
+        # changes the output.  (See https://bugs.python.org/issue34093.)
+        code = marshal.loads(data)
+        self.assertEqual(code, expected)
 
     def test_module(self):
-        names = [
-            '__hello__',
-            '__hello_alias__',
-            '__hello_only__',
-            '__phello__.__init__',
-            '__phello__.spam',
-            '__phello__.ham.__init__',
-            '__phello__.ham.eggs',
-        ]
-        for name in names:
+        modules = {
+            '__hello__': None,
+            '__phello__.__init__': None,
+            '__phello__.spam': None,
+            '__phello__.ham.__init__': None,
+            '__phello__.ham.eggs': None,
+            '__hello_alias__': '__hello__',
+            }
+        for name, source in modules.items():
             with self.subTest(name):
                 spec = self.find(name)
-                self.check(spec, name)
-                self.assertEqual(spec.submodule_search_locations, None)
+                self.check_basic(spec, name)
+                self.check_data(spec, source)
 
     def test_package(self):
-        names = [
-            '__phello__',
-            '__phello__.ham',
-            '__phello_alias__',
-        ]
-        for name in names:
+        modules = {
+            '__phello__': None,
+            '__phello__.ham': None,
+            '__phello_alias__': '__hello__',
+        }
+        for name, source in modules.items():
             with self.subTest(name):
                 spec = self.find(name)
-                self.check(spec, name)
-                self.assertEqual(spec.submodule_search_locations, [])
+                self.check_basic(spec, name, ispkg=True)
+                self.check_search_location(spec, source)
+                self.check_data(spec, source, ispkg=True)
+
+    def test_frozen_only(self):
+        name = '__hello_only__'
+        source = os.path.join(REPO_ROOT, 'Tools', 'freeze', 'flag.py')
+        spec = self.find(name)
+        self.check_basic(spec, name)
+        self.check_data(spec, source)
 
     # These are covered by test_module() and test_package().
     test_module_in_package = None
diff --git a/Lib/test/test_importlib/frozen/test_loader.py b/Lib/test/test_importlib/frozen/test_loader.py
index cfa5e5b..992dcef 100644
--- a/Lib/test/test_importlib/frozen/test_loader.py
+++ b/Lib/test/test_importlib/frozen/test_loader.py
@@ -4,7 +4,9 @@ from .. import util
 machinery = util.import_importlib('importlib.machinery')
 
 from test.support import captured_stdout, import_helper
+import _imp
 import contextlib
+import marshal
 import types
 import unittest
 import warnings
@@ -33,11 +35,14 @@ class ExecModuleTests(abc.LoaderTests):
     def exec_module(self, name):
         with import_helper.frozen_modules():
             is_package = self.machinery.FrozenImporter.is_package(name)
+            code = _imp.get_frozen_object(name)
+        data = marshal.dumps(code)
         spec = self.machinery.ModuleSpec(
             name,
             self.machinery.FrozenImporter,
             origin='frozen',
             is_package=is_package,
+            loader_state=data,
         )
         module = types.ModuleType(name)
         module.__spec__ = spec
@@ -61,6 +66,7 @@ class ExecModuleTests(abc.LoaderTests):
             self.assertEqual(getattr(module, attr), value)
         self.assertEqual(output, 'Hello world!\n')
         self.assertTrue(hasattr(module, '__spec__'))
+        self.assertIsNone(module.__spec__.loader_state)
 
     def test_package(self):
         name = '__phello__'
@@ -73,6 +79,7 @@ class ExecModuleTests(abc.LoaderTests):
                                  name=name, attr=attr, given=attr_value,
                                  expected=value))
         self.assertEqual(output, 'Hello world!\n')
+        self.assertIsNone(module.__spec__.loader_state)
 
     def test_lacking_parent(self):
         name = '__phello__.spam'