1 files changed, 38 insertions, 18 deletions

diff --git a/Python/import.c b/Python/import.c
index 010e821..80fe5b7 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -2211,11 +2211,45 @@ ensure_fromlist(PyObject *mod, PyObject *fromlist, char *buf, int buflen,
 	/* NOTREACHED */
 }
 
+static int
+add_submodule(PyObject *mod, PyObject *submod, char *fullname, char *subname,
+	      PyObject *modules)
+{
+	if (mod == Py_None)
+		return 1;
+	/* Irrespective of the success of this load, make a
+	   reference to it in the parent package module.  A copy gets
+	   saved in the modules dictionary under the full name, so get a
+	   reference from there, if need be.  (The exception is when the
+	   load failed with a SyntaxError -- then there's no trace in
+	   sys.modules.  In that case, of course, do nothing extra.) */
+	if (submod == NULL) {
+		submod = PyDict_GetItemString(modules, fullname);
+		if (submod == NULL)
+			return 1;
+	}
+	if (PyModule_Check(mod)) {
+		/* We can't use setattr here since it can give a
+		 * spurious warning if the submodule name shadows a
+		 * builtin name */
+		PyObject *dict = PyModule_GetDict(mod);
+		if (!dict)
+			return 0;
+		if (PyDict_SetItemString(dict, subname, submod) < 0)
+			return 0;
+	}
+	else {
+		if (PyObject_SetAttrString(mod, subname, submod) < 0)
+			return 0;
+	}
+	return 1;
+}
+
 static PyObject *
 import_submodule(PyObject *mod, char *subname, char *fullname)
 {
 	PyObject *modules = PyImport_GetModuleDict();
-	PyObject *m, *res = NULL;
+	PyObject *m = NULL;
 
 	/* Require:
 	   if mod == None: subname == fullname
@@ -2257,23 +2291,9 @@ import_submodule(PyObject *mod, char *subname, char *fullname)
 		Py_XDECREF(loader);
 		if (fp)
 			fclose(fp);
-		if (mod != Py_None) {
-			/* Irrespective of the success of this load, make a
-			   reference to it in the parent package module.
-			   A copy gets saved in the modules dictionary
-			   under the full name, so get a reference from
-			   there, if need be.  (The exception is when
-			   the load failed with a SyntaxError -- then
-			   there's no trace in sys.modules.  In that case,
-			   of course, do nothing extra.) */
-			res = m;
-			if (res == NULL)
-				res = PyDict_GetItemString(modules, fullname);
-			if (res != NULL &&
-			    PyObject_SetAttrString(mod, subname, res) < 0) {
-				Py_XDECREF(m);
-				m = NULL;
-			}
+		if (!add_submodule(mod, m, fullname, subname, modules)) {
+			Py_XDECREF(m);
+			m = NULL;
 		}
 	}