gh-87533: Expand pickle importing to support non-package C-modules

There have been recurring issue with PyModule_Create modules in PyTorch; when trying to serialize attributes of these C-modules, pickle fails to import the C-module because it is not a package

This is the current issue that brought this to my attention: pytorch/pytorch#126154

The existing hack to this issue has been to insert the C-module into sys.modules in order to enable pickle to find them: https://github.com/pytorch/pytorch/pull/38136/files#diff-d7e90d0f94b43db763b44fba679a5c1b4cabe3668aaf34f2aee07de8e2d1b2faR524-R528

Instead of relying on this hack, we can change `pickle`'s  approach to loading, which is currently equivalent to `import package.c_module`; instead, we could do `from package import c_module`, which 1) does not care if `c_module` is a package or not 2) is fully backward compatible with the previous approach and 3) slots in nicely to the `fromlist` parameter of `__import__`, which we are already using to load modules in `pickle`

Author: Matthew Hoffman

Lib/pickle.py

@@ -345,6 +345,24 @@ def whichmodule(obj, name):
             pass
     return '__main__'
 
+def import_module_from_string(module_name):
+    """Import a module from a string.
+
+    The last module in the dot-delimited module name is returned. The last
+    module need not be a package (see bpo-43367).
+
+    >>> import_module_from_string('pickle')
+    <module 'pickle' from '...'>
+    >>> import_module_from_string('collections.abc)
+    <module 'collections.abc' from '...'>
+    """
+    module_name, *fromlist = module_name.rsplit(".", 1)
+    module = __import__(module_name, fromlist=fromlist, level=0)
+    if fromlist:
+        assert len(fromlist) == 1
+        module = getattr(module, fromlist[0])
+    return module
+
 def encode_long(x):
     r"""Encode a long to a two's complement little-endian binary string.
     Note that 0 is a special case, returning an empty string, to save a
@@ -364,7 +382,6 @@ def encode_long(x):
     b'\x80'
     >>> encode_long(127)
     b'\x7f'
-    >>>
     """
     if x == 0:
         return b''
@@ -1049,7 +1066,6 @@ def save_frozenset(self, obj):
 
     def save_global(self, obj, name=None):
         write = self.write
-        memo = self.memo
 
         if name is None:
             name = getattr(obj, '__qualname__', None)
@@ -1058,8 +1074,7 @@ def save_global(self, obj, name=None):
 
         module_name = whichmodule(obj, name)
         try:
-            __import__(module_name, level=0)
-            module = sys.modules[module_name]
+            module = import_module_from_string(module_name)
             obj2, parent = _getattribute(module, name)
         except (ImportError, KeyError, AttributeError):
             raise PicklingError(
@@ -1565,17 +1580,18 @@ def get_extension(self, code):
 
     def find_class(self, module, name):
         # Subclasses may override this.
-        sys.audit('pickle.find_class', module, name)
+        module_name = module
+        sys.audit('pickle.find_class', module_name, name)
         if self.proto < 3 and self.fix_imports:
-            if (module, name) in _compat_pickle.NAME_MAPPING:
-                module, name = _compat_pickle.NAME_MAPPING[(module, name)]
-            elif module in _compat_pickle.IMPORT_MAPPING:
-                module = _compat_pickle.IMPORT_MAPPING[module]
-        __import__(module, level=0)
+            if (module_name, name) in _compat_pickle.NAME_MAPPING:
+                module_name, name = _compat_pickle.NAME_MAPPING[(module_name, name)]
+            elif module_name in _compat_pickle.IMPORT_MAPPING:
+                module_name = _compat_pickle.IMPORT_MAPPING[module_name]
+        module = import_module_from_string(module_name)
         if self.proto >= 4:
-            return _getattribute(sys.modules[module], name)[0]
+            return _getattribute(module, name)[0]
         else:
-            return getattr(sys.modules[module], name)
+            return getattr(module, name)
 
     def load_reduce(self):
         stack = self.stack

Lib/test/pickletester.py

@@ -2845,6 +2845,12 @@ class Subclass(tuple):
             class Nested(str):
                 pass
 
+        c_module = types.ModuleType("c_module")
+        def c_function():
+            return None
+        c_function.__module__ = f"{__name__}.{c_module.__name__}"
+        c_module.c_function = c_function
+
         c_methods = (
             # bound built-in method
             ("abcd".index, ("c",)),
@@ -2867,6 +2873,8 @@ class Nested(str):
             (Subclass.count, (Subclass([1,2,2]), 2)),
             (Subclass.Nested("sweet").count, ("e",)),
             (Subclass.Nested.count, (Subclass.Nested("sweet"), "e")),
+            # bpo-43367: pickling C-module attributes
+            (c_module.c_function, ()),
         )
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             for method, args in c_methods:

Modules/_pickle.c

@@ -3584,6 +3584,63 @@ fix_imports(PickleState *st, PyObject **module_name, PyObject **global_name)
     return 0;
 }
 
+PyObject*
+import_module_from_string(PyObject *module_name)
+{
+    PyObject *split_module_name = PyUnicode_RSplit(module_name,
+                                                   PyUnicode_FromString("."),
+                                                   1);
+    if (split_module_name == NULL) {
+        PyErr_Format(PyExc_RuntimeError,
+                     "Failed to split module name %R", module_name);
+        return NULL;
+    }
+
+    module_name = PySequence_Fast_GET_ITEM(split_module_name, 0);
+    if (module_name == NULL) {
+        PyErr_Format(PyExc_RuntimeError, "Failed to get module name from %R",
+                     split_module_name);
+        return NULL;
+    }
+
+    PyObject *fromlist = PySequence_GetSlice(split_module_name, 1,
+                                             PySequence_Fast_GET_SIZE(split_module_name));
+    if (fromlist == NULL) {
+        PyErr_Format(PyExc_RuntimeError, "Failed to get fromlist from %R",
+                     split_module_name);
+        return NULL;
+    }
+
+    PyObject *module = PyImport_ImportModuleLevelObject(module_name, NULL,
+                                                        NULL, fromlist, 0);
+    if (module == NULL) {
+        PyErr_Format(PyExc_ModuleNotFoundError, "Import of module %R failed",
+                     module_name);
+        return NULL;
+    }
+
+    Py_ssize_t fromlist_size = PySequence_Fast_GET_SIZE(fromlist);
+    if (fromlist_size > 0) {
+        assert(fromlist_size == 1);
+        PyObject *submodule_name = PySequence_Fast_GET_ITEM(fromlist, 0);
+        if (submodule_name == NULL) {
+            PyErr_Format(PyExc_RuntimeError, "Failed to get submodule name from %R",
+                         fromlist);
+            return NULL;
+        }
+        module = PyObject_GetAttr(module, submodule_name);
+        if (module == NULL) {
+            PyErr_Format(PyExc_ModuleNotFoundError,
+                         "Attribute lookup %R on %R failed", submodule_name,
+                         module_name);
+            return NULL;
+        }
+    }
+
+    return module;
+}
+
+
 static int
 save_global(PickleState *st, PicklerObject *self, PyObject *obj,
             PyObject *name)
@@ -3619,21 +3676,11 @@ save_global(PickleState *st, PicklerObject *self, PyObject *obj,
     if (module_name == NULL)
         goto error;
 
-    /* XXX: Change to use the import C API directly with level=0 to disallow
-       relative imports.
-
-       XXX: PyImport_ImportModuleLevel could be used. However, this bypasses
-       builtins.__import__. Therefore, _pickle, unlike pickle.py, will ignore
-       custom import functions (IMHO, this would be a nice security
-       feature). The import C API would need to be extended to support the
-       extra parameters of __import__ to fix that. */
-    module = PyImport_Import(module_name);
+    module = import_module_from_string(module_name);
     if (module == NULL) {
-        PyErr_Format(st->PicklingError,
-                     "Can't pickle %R: import of module %R failed",
-                     obj, module_name);
         goto error;
     }
+
     lastname = Py_NewRef(PyList_GET_ITEM(dotted_path,
                          PyList_GET_SIZE(dotted_path) - 1));
     cls = get_deep_attribute(module, dotted_path, &parent);
@@ -7036,7 +7083,7 @@ _pickle_Unpickler_find_class_impl(UnpicklerObject *self, PyTypeObject *cls,
      * we don't use PyImport_GetModule here, because it can return partially-
      * initialised modules, which then cause the getattribute to fail.
      */
-    module = PyImport_Import(module_name);
+    module = import_module_from_string(module_name);
     if (module == NULL) {
         return NULL;
     }