Provide __module__ attributes for functions defined in C and Python.

__module__ is the string name of the module the function was defined
in, just like __module__ of classes.  In some cases, particularly for
C functions, the __module__ may be None.

Change PyCFunction_New() from a function to a macro, but keep an
unused copy of the function around so that we don't change the binary
API.

Change pickle's save_global() to use whichmodule() if __module__ is
None, but add the __module__ logic to whichmodule() since it might be
used outside of pickle.

2 files changed, 68 insertions, 6 deletions

diff --git a/Objects/funcobject.c b/Objects/funcobject.c
index 6154d99..b162fdf 100644
--- a/Objects/funcobject.c
+++ b/Objects/funcobject.c
@@ -14,6 +14,7 @@ PyFunction_New(PyObject *code, PyObject *globals)
 	if (op != NULL) {
 		PyObject *doc;
 		PyObject *consts;
+		PyObject *module;
 		op->func_weakreflist = NULL;
 		Py_INCREF(code);
 		op->func_code = code;
@@ -34,6 +35,16 @@ PyFunction_New(PyObject *code, PyObject *globals)
 		Py_INCREF(doc);
 		op->func_doc = doc;
 		op->func_dict = NULL;
+		op->func_module = NULL;
+
+		/* __module__: If module name is in globals, use it.
+		   Otherwise, use None.
+		*/
+		module = PyDict_GetItemString(globals, "__name__");
+		if (module) {
+		    Py_INCREF(module);
+		    op->func_module = module;
+		}
 	}
 	else
 		return NULL;
@@ -62,6 +73,16 @@ PyFunction_GetGlobals(PyObject *op)
 }
 
 PyObject *
+PyFunction_GetModule(PyObject *op)
+{
+	if (!PyFunction_Check(op)) {
+		PyErr_BadInternalCall();
+		return NULL;
+	}
+	return ((PyFunctionObject *) op) -> func_module;
+}
+
+PyObject *
 PyFunction_GetDefaults(PyObject *op)
 {
 	if (!PyFunction_Check(op)) {
@@ -138,6 +159,7 @@ static PyMemberDef func_memberlist[] = {
 	 RESTRICTED|READONLY},
         {"func_name",     T_OBJECT,     OFF(func_name),         READONLY},
         {"__name__",      T_OBJECT,     OFF(func_name),         READONLY},
+	{"__module__",    T_OBJECT,     OFF(func_module),       READONLY},
         {NULL}  /* Sentinel */
 };
 
@@ -373,6 +395,7 @@ func_dealloc(PyFunctionObject *op)
 		PyObject_ClearWeakRefs((PyObject *) op);
 	Py_DECREF(op->func_code);
 	Py_DECREF(op->func_globals);
+	Py_XDECREF(op->func_module);
 	Py_DECREF(op->func_name);
 	Py_XDECREF(op->func_defaults);
 	Py_XDECREF(op->func_doc);
@@ -405,6 +428,11 @@ func_traverse(PyFunctionObject *f, visitproc visit, void *arg)
 		if (err)
 			return err;
 	}
+	if (f->func_module) {
+		err = visit(f->func_module, arg);
+		if (err)
+			return err;
+	}
 	if (f->func_defaults) {
 		err = visit(f->func_defaults, arg);
 		if (err)
diff --git a/Objects/methodobject.c b/Objects/methodobject.c
index a1f325d..aa64ff1 100644
--- a/Objects/methodobject.c
+++ b/Objects/methodobject.c
@@ -2,11 +2,12 @@
 /* Method object implementation */
 
 #include "Python.h"
+#include "structmember.h"
 
 static PyCFunctionObject *free_list = NULL;
 
 PyObject *
-PyCFunction_New(PyMethodDef *ml, PyObject *self)
+PyCFunction_NewEx(PyMethodDef *ml, PyObject *self, PyObject *module)
 {
 	PyCFunctionObject *op;
 	op = free_list;
@@ -22,6 +23,8 @@ PyCFunction_New(PyMethodDef *ml, PyObject *self)
 	op->m_ml = ml;
 	Py_XINCREF(self);
 	op->m_self = self;
+	Py_XINCREF(module);
+	op->m_module = module;
 	_PyObject_GC_TRACK(op);
 	return (PyObject *)op;
 }
@@ -121,6 +124,7 @@ meth_dealloc(PyCFunctionObject *m)
 {
 	_PyObject_GC_UNTRACK(m);
 	Py_XDECREF(m->m_self);
+	Py_XDECREF(m->m_module);
 	m->m_self = (PyObject *)free_list;
 	free_list = m;
 }
@@ -145,10 +149,18 @@ meth_get__name__(PyCFunctionObject *m, void *closure)
 static int
 meth_traverse(PyCFunctionObject *m, visitproc visit, void *arg)
 {
-	if (m->m_self != NULL)
-		return visit(m->m_self, arg);
-	else
-		return 0;
+	int err;
+	if (m->m_self != NULL) {
+		err = visit(m->m_self, arg);
+		if (err)
+			return err;
+	}
+	if (m->m_module != NULL) {
+		err = visit(m->m_module, arg);
+		if (err)
+			return err;
+	}
+	return 0;
 }
 
 static PyObject *
@@ -174,6 +186,13 @@ static PyGetSetDef meth_getsets [] = {
 	{0}
 };
 
+#define OFF(x) offsetof(PyCFunctionObject, x)
+
+static PyMemberDef meth_members[] = {
+	{"__module__",    T_OBJECT,     OFF(m_module),       READONLY},
+	{NULL}
+};
+
 static PyObject *
 meth_repr(PyCFunctionObject *m)
 {
@@ -250,7 +269,7 @@ PyTypeObject PyCFunction_Type = {
 	0,					/* tp_iter */
 	0,					/* tp_iternext */
 	0,					/* tp_methods */
-	0,					/* tp_members */
+	meth_members,				/* tp_members */
 	meth_getsets,				/* tp_getset */
 	0,					/* tp_base */
 	0,					/* tp_dict */
@@ -308,6 +327,7 @@ Py_FindMethodInChain(PyMethodChain *chain, PyObject *self, char *name)
 		for (; ml->ml_name != NULL; ml++) {
 			if (name[0] == ml->ml_name[0] &&
 			    strcmp(name+1, ml->ml_name+1) == 0)
+				/* XXX */
 				return PyCFunction_New(ml, self);
 		}
 		chain = chain->link;
@@ -338,3 +358,17 @@ PyCFunction_Fini(void)
 		PyObject_GC_Del(v);
 	}
 }
+
+/* PyCFunction_New() is now just a macro that calls PyCFunction_NewEx(),
+   but it's part of the API so we need to keep a function around that
+   existing C extensions can call.
+*/
+   
+#undef PyCFunction_New
+PyAPI_FUNC(PyObject *) PyCFunction_New(PyMethodDef *, PyObject *);
+
+PyObject *
+PyCFunction_New(PyMethodDef *ml, PyObject *self)
+{
+	return PyCFunction_NewEx(ml, self, NULL);
+}