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

1 files changed, 57 insertions, 10 deletions

diff --git a/Python/clinic/import.c.h b/Python/clinic/import.c.h
index 438a348..0973883 100644
--- a/Python/clinic/import.c.h
+++ b/Python/clinic/import.c.h
@@ -169,33 +169,80 @@ exit:
     return return_value;
 }
 
-PyDoc_STRVAR(_imp_get_frozen_object__doc__,
-"get_frozen_object($module, name, /)\n"
+PyDoc_STRVAR(_imp_find_frozen__doc__,
+"find_frozen($module, name, /)\n"
 "--\n"
 "\n"
-"Create a code object for a frozen module.");
+"Return info about the corresponding frozen module (if there is one) or None.\n"
+"\n"
+"The returned info (a 2-tuple):\n"
+"\n"
+" * data         the raw marshalled bytes\n"
+" * is_package   whether or not it is a package");
 
-#define _IMP_GET_FROZEN_OBJECT_METHODDEF    \
-    {"get_frozen_object", (PyCFunction)_imp_get_frozen_object, METH_O, _imp_get_frozen_object__doc__},
+#define _IMP_FIND_FROZEN_METHODDEF    \
+    {"find_frozen", (PyCFunction)_imp_find_frozen, METH_O, _imp_find_frozen__doc__},
 
 static PyObject *
-_imp_get_frozen_object_impl(PyObject *module, PyObject *name);
+_imp_find_frozen_impl(PyObject *module, PyObject *name);
 
 static PyObject *
-_imp_get_frozen_object(PyObject *module, PyObject *arg)
+_imp_find_frozen(PyObject *module, PyObject *arg)
 {
     PyObject *return_value = NULL;
     PyObject *name;
 
     if (!PyUnicode_Check(arg)) {
-        _PyArg_BadArgument("get_frozen_object", "argument", "str", arg);
+        _PyArg_BadArgument("find_frozen", "argument", "str", arg);
         goto exit;
     }
     if (PyUnicode_READY(arg) == -1) {
         goto exit;
     }
     name = arg;
-    return_value = _imp_get_frozen_object_impl(module, name);
+    return_value = _imp_find_frozen_impl(module, name);
+
+exit:
+    return return_value;
+}
+
+PyDoc_STRVAR(_imp_get_frozen_object__doc__,
+"get_frozen_object($module, name, data=None, /)\n"
+"--\n"
+"\n"
+"Create a code object for a frozen module.");
+
+#define _IMP_GET_FROZEN_OBJECT_METHODDEF    \
+    {"get_frozen_object", (PyCFunction)(void(*)(void))_imp_get_frozen_object, METH_FASTCALL, _imp_get_frozen_object__doc__},
+
+static PyObject *
+_imp_get_frozen_object_impl(PyObject *module, PyObject *name,
+                            PyObject *dataobj);
+
+static PyObject *
+_imp_get_frozen_object(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    PyObject *name;
+    PyObject *dataobj = Py_None;
+
+    if (!_PyArg_CheckPositional("get_frozen_object", nargs, 1, 2)) {
+        goto exit;
+    }
+    if (!PyUnicode_Check(args[0])) {
+        _PyArg_BadArgument("get_frozen_object", "argument 1", "str", args[0]);
+        goto exit;
+    }
+    if (PyUnicode_READY(args[0]) == -1) {
+        goto exit;
+    }
+    name = args[0];
+    if (nargs < 2) {
+        goto skip_optional;
+    }
+    dataobj = args[1];
+skip_optional:
+    return_value = _imp_get_frozen_object_impl(module, name, dataobj);
 
 exit:
     return return_value;
@@ -498,4 +545,4 @@ exit:
 #ifndef _IMP_EXEC_DYNAMIC_METHODDEF
     #define _IMP_EXEC_DYNAMIC_METHODDEF
 #endif /* !defined(_IMP_EXEC_DYNAMIC_METHODDEF) */
-/*[clinic end generated code: output=96038c277119d6e3 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=a31e1c00653359ff input=a9049054013a1b77]*/