Manual merge of PEP 366 implementation from trunk (the automatic merge choked on the PyString->PyUnicode changes)

1 files changed, 86 insertions, 26 deletions

diff --git a/Python/import.c b/Python/import.c
index b7a9752..9db8fb0 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -2016,7 +2016,8 @@ get_parent(PyObject *globals, char *buf, Py_ssize_t *p_buflen, int level)
 {
 	static PyObject *namestr = NULL;
 	static PyObject *pathstr = NULL;
-	PyObject *modname, *modpath, *modules, *parent;
+	static PyObject *pkgstr = NULL;
+	PyObject *pkgname, *modname, *modpath, *modules, *parent;
 
 	if (globals == NULL || !PyDict_Check(globals) || !level)
 		return Py_None;
@@ -2031,44 +2032,103 @@ get_parent(PyObject *globals, char *buf, Py_ssize_t *p_buflen, int level)
 		if (pathstr == NULL)
 			return NULL;
 	}
+	if (pkgstr == NULL) {
+		pkgstr = PyUnicode_InternFromString("__package__");
+		if (pkgstr == NULL)
+			return NULL;
+	}
 
 	*buf = '\0';
 	*p_buflen = 0;
-	modname = PyDict_GetItem(globals, namestr);
-	if (modname == NULL || !PyUnicode_Check(modname))
-		return Py_None;
+	pkgname = PyDict_GetItem(globals, pkgstr);
 
-	modpath = PyDict_GetItem(globals, pathstr);
-	if (modpath != NULL) {
-		Py_ssize_t len = PyUnicode_GET_SIZE(modname);
-		if (len > MAXPATHLEN) {
+	if ((pkgname != NULL) && (pkgname != Py_None)) {
+		/* __package__ is set, so use it */
+		Py_ssize_t len;
+		if (!PyUnicode_Check(pkgname)) {
 			PyErr_SetString(PyExc_ValueError,
-					"Module name too long");
+					"__package__ set to non-string");
 			return NULL;
 		}
-		strcpy(buf, PyUnicode_AsString(modname));
-	}
-	else {
-		char *start = PyUnicode_AsString(modname);
-		char *lastdot = strrchr(start, '.');
-		size_t len;
-		if (lastdot == NULL && level > 0) {
+		len = PyUnicode_GET_SIZE(pkgname);
+		if (len == 0) {
+			if (level > 0) {
+				PyErr_SetString(PyExc_ValueError,
+					"Attempted relative import in non-package");
+				return NULL;
+			}
+			return Py_None;
+		}
+		if (len > MAXPATHLEN) {
 			PyErr_SetString(PyExc_ValueError,
-				"Attempted relative import in non-package");
+					"Package name too long");
 			return NULL;
 		}
-		if (lastdot == NULL)
+		strcpy(buf, PyUnicode_AsString(pkgname));
+	} else {
+		/* __package__ not set, so figure it out and set it */
+		modname = PyDict_GetItem(globals, namestr);
+		if (modname == NULL || !PyUnicode_Check(modname))
 			return Py_None;
-		len = lastdot - start;
-		if (len >= MAXPATHLEN) {
-			PyErr_SetString(PyExc_ValueError,
-					"Module name too long");
-			return NULL;
+	
+		modpath = PyDict_GetItem(globals, pathstr);
+		if (modpath != NULL) {
+			/* __path__ is set, so modname is already the package name */
+			Py_ssize_t len = PyUnicode_GET_SIZE(modname);
+			int error;
+			if (len > MAXPATHLEN) {
+				PyErr_SetString(PyExc_ValueError,
+						"Module name too long");
+				return NULL;
+			}
+			strcpy(buf, PyUnicode_AsString(modname));
+			error = PyDict_SetItem(globals, pkgstr, modname);
+			if (error) {
+				PyErr_SetString(PyExc_ValueError,
+						"Could not set __package__");
+				return NULL;
+			}
+		} else {
+			/* Normal module, so work out the package name if any */
+			char *start = PyUnicode_AsString(modname);
+			char *lastdot = strrchr(start, '.');
+			size_t len;
+			int error;
+			if (lastdot == NULL && level > 0) {
+				PyErr_SetString(PyExc_ValueError,
+					"Attempted relative import in non-package");
+				return NULL;
+			}
+			if (lastdot == NULL) {
+				error = PyDict_SetItem(globals, pkgstr, Py_None);
+				if (error) {
+					PyErr_SetString(PyExc_ValueError,
+						"Could not set __package__");
+					return NULL;
+				}
+				return Py_None;
+			}
+			len = lastdot - start;
+			if (len >= MAXPATHLEN) {
+				PyErr_SetString(PyExc_ValueError,
+						"Module name too long");
+				return NULL;
+			}
+			strncpy(buf, start, len);
+			buf[len] = '\0';
+			pkgname = PyUnicode_FromString(buf);
+			if (pkgname == NULL) {
+				return NULL;
+			}
+			error = PyDict_SetItem(globals, pkgstr, pkgname);
+			Py_DECREF(pkgname);
+			if (error) {
+				PyErr_SetString(PyExc_ValueError,
+						"Could not set __package__");
+				return NULL;
+			}
 		}
-		strncpy(buf, start, len);
-		buf[len] = '\0';
 	}
-
 	while (--level > 0) {
 		char *dot = strrchr(buf, '.');
 		if (dot == NULL) {