Add origname to spec.loader_state.

python · ericsnowcurrently · Oct 5, 2021 · Sep 28, 2021 · Sep 28, 2021 · Sep 29, 2021
commit 905d5f5747a8f178cc0591cab51ef70f1008ce79
@@ -829,12 +829,13 @@ def find_spec(cls, fullname, path=None, target=None):
         info = _call_with_frames_removed(_imp.find_frozen, fullname)
         if info is None:
             return None
-        data, ispkg = info
+        data, ispkg, origname = info
         spec = spec_from_loader(fullname, cls,
                                 origin=cls._ORIGIN,
                                 is_package=ispkg)
         spec.loader_state = type(sys.implementation)(
             data=data,
+            origname=origname,
         )
         return spec
 

@@ -9,7 +9,15 @@
 import unittest
 import warnings
 
-from test.support import import_helper, REPO_ROOT
+from test.support import import_helper, REPO_ROOT, STDLIB_DIR
+
+
+def resolve_stdlib_file(name, ispkg=False):
+    assert name
+    if ispkg:
+        return os.path.join(STDLIB_DIR, *name.split('.'), '__init__.py')
+    else:
+        return os.path.join(STDLIB_DIR, *name.split('.')) + '.py'
 
 
 class FindSpecTests(abc.FinderTests):
@@ -32,7 +40,11 @@ def check_basic(self, spec, name, ispkg=False):
             self.assertIsNone(spec.submodule_search_locations)
         self.assertIsNotNone(spec.loader_state)
 
-    def check_loader_state(self, spec):
+    def check_loader_state(self, spec, origname=None, filename=None):
+        if not filename:
+            if not origname:
+                origname = spec.name
+
         actual = dict(vars(spec.loader_state))
 
         # Check the code object used to import the frozen module.
@@ -47,7 +59,9 @@ def check_loader_state(self, spec):
         self.assertEqual(code, expected)
 
         # Check the rest of spec.loader_state.
-        expected = dict()
+        expected = dict(
+            origname=origname,
+        )
         self.assertDictEqual(actual, expected)
 
     def check_search_locations(self, spec):
@@ -75,46 +89,50 @@ def test_module(self):
             with self.subTest(f'{name} -> {origname}'):
                 spec = self.find(name)
                 self.check_basic(spec, name)
-                self.check_loader_state(spec)
+                self.check_loader_state(spec, origname)
         modules = [
             '__phello__.__init__',
             '__phello__.ham.__init__',
         ]
         for name in modules:
             origname = name.rpartition('.')[0]
+            filename = resolve_stdlib_file(name)
             with self.subTest(f'{name} -> {origname}'):
                 spec = self.find(name)
                 self.check_basic(spec, name)
-                self.check_loader_state(spec)
+                self.check_loader_state(spec, origname, filename)
         modules = {
             '__hello_only__': ('Tools', 'freeze', 'flag.py'),
         }
         for name, path in modules.items():
+            origname = None
             filename = os.path.join(REPO_ROOT, *path)
             with self.subTest(f'{name} -> {filename}'):
                 spec = self.find(name)
                 self.check_basic(spec, name)
-                self.check_loader_state(spec)
+                self.check_loader_state(spec, origname, filename)
 
     def test_package(self):
         packages = [
             '__phello__',
             '__phello__.ham',
         ]
         for name in packages:
+            filename = resolve_stdlib_file(name, ispkg=True)
             with self.subTest(f'{name} -> {name}'):
                 spec = self.find(name)
                 self.check_basic(spec, name, ispkg=True)
-                self.check_loader_state(spec)
+                self.check_loader_state(spec, name, filename)
                 self.check_search_locations(spec)
         packages = {
             '__phello_alias__': '__hello__',
         }
         for name, origname in packages.items():
+            filename = resolve_stdlib_file(origname, ispkg=False)
             with self.subTest(f'{name} -> {origname}'):
                 spec = self.find(name)
                 self.check_basic(spec, name, ispkg=True)
-                self.check_loader_state(spec)
+                self.check_loader_state(spec, origname, filename)
                 self.check_search_locations(spec)
 
     # These are covered by test_module() and test_package().

@@ -30,19 +30,29 @@ def fresh(name, *, oldapi=False):
                 yield
 
 
+def get_origname(name):
+    if name.endswith('.__init__'):
+        return name.rpartition('.')[0]
+    else:
+        # This is good enough for now.
+        return name
+
+
 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=types.SimpleNamespace(data=data),
+            loader_state=types.SimpleNamespace(
+                data=marshal.dumps(code),
+                origname=get_origname(name),
+            ),
         )
         module = types.ModuleType(name)
         module.__spec__ = spec

@@ -2039,11 +2039,14 @@ The returned info (a 2-tuple):
 
  * data         the raw marshalled bytes
  * is_package   whether or not it is a package
+ * origname     the originally frozen module's name, or None if not
+                a stdlib module (this will usually be the same as
+                the module's current name)
 [clinic start generated code]*/
 
 static PyObject *
 _imp_find_frozen_impl(PyObject *module, PyObject *name)
-/*[clinic end generated code: output=3fd17da90d417e4e input=4e52b3ac95f6d7ab]*/
+/*[clinic end generated code: output=3fd17da90d417e4e input=6aa7b9078a89280a]*/
 {
     struct frozen_info info;
     frozen_status status = find_frozen(name, &info);
@@ -2057,12 +2060,25 @@ _imp_find_frozen_impl(PyObject *module, PyObject *name)
         set_frozen_error(status, name);
         return NULL;
     }
+
     PyObject *data = PyBytes_FromStringAndSize(info.data, info.size);
     if (data == NULL) {
         return NULL;
     }
-    PyObject *result = PyTuple_Pack(2, data,
-                                    info.is_package ? Py_True : Py_False);
+
+    PyObject *origname = NULL;
+    if (info.origname != NULL && info.origname[0] != '\0') {
+        origname = PyUnicode_FromString(info.origname);
+        if (origname == NULL) {
+            Py_DECREF(data);
+            return NULL;
+        }
+    }
+
+    PyObject *result = PyTuple_Pack(3, data,
+                                    info.is_package ? Py_True : Py_False,
+                                    origname ? origname : Py_None);
+    Py_XDECREF(origname);
     Py_DECREF(data);
     return result;
 }