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

Author: ericsnowcurrently

Lib/importlib/_bootstrap.py

@@ -826,10 +826,15 @@ def module_repr(m):
 
     @classmethod
     def find_spec(cls, fullname, path=None, target=None):
-        if _imp.is_frozen(fullname):
-            return spec_from_loader(fullname, cls, origin=cls._ORIGIN)
-        else:
+        info = _call_with_frames_removed(_imp.find_frozen, fullname)
+        if info is None:
             return None
+        data, ispkg = info
+        spec = spec_from_loader(fullname, cls,
+                                origin=cls._ORIGIN,
+                                is_package=ispkg)
+        spec.loader_state = data
+        return spec
 
     @classmethod
     def find_module(cls, fullname, path=None):
@@ -849,11 +854,22 @@ def create_module(spec):
 
     @staticmethod
     def exec_module(module):
-        name = module.__spec__.name
-        if not _imp.is_frozen(name):
-            raise ImportError('{!r} is not a frozen module'.format(name),
-                              name=name)
-        code = _call_with_frames_removed(_imp.get_frozen_object, name)
+        spec = module.__spec__
+        name = spec.name
+        try:
+            data = spec.loader_state
+        except AttributeError:
+            if not _imp.is_frozen(name):
+                raise ImportError('{!r} is not a frozen module'.format(name),
+                                  name=name)
+            data = None
+        else:
+            # We clear the extra data we got from the finder, to save memory.
+            # Note that if this method is called again (e.g. by
+            # importlib.reload()) then _imp.get_frozen_object() will notice
+            # no data was provided and will look it up.
+            spec.loader_state = None
+        code = _call_with_frames_removed(_imp.get_frozen_object, name, data)
         exec(code, module.__dict__)
 
     @classmethod

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 @@ def find(self, name, **kwargs):
         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

Lib/test/test_importlib/frozen/test_loader.py

@@ -4,7 +4,9 @@
 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 @@ def test_module(self):
             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 @@ def test_package(self):
                                  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'

Misc/NEWS.d/next/Core and Builtins/2021-09-29-12-02-39.bpo-45324.BTQElX.rst

@@ -0,0 +1,3 @@
+In FrozenImporter.find_spec(), we now preserve the information needed in
+exec_module() to load the module.  This change mostly impacts internal
+details, rather than changing the importer's behavior.