Patch #1591665: implement the __dir__() special function lookup in PyObject_Dir.

1 files changed, 144 insertions, 126 deletions

diff --git a/Objects/object.c b/Objects/object.c
index 065abdc..e2d1b13 100644
--- a/Objects/object.c
+++ b/Objects/object.c
@@ -1284,6 +1284,8 @@ PyCallable_Check(PyObject *x)
 	return x->ob_type->tp_call != NULL;
 }
 
+/* ------------------------- PyObject_Dir() helpers ------------------------- */
+
 /* Helper for PyObject_Dir.
    Merge the __dict__ of aclass into dict, and recursively also all
    the __dict__s of aclass's base classes.  The order of merging isn't
@@ -1343,158 +1345,174 @@ merge_class_dict(PyObject* dict, PyObject* aclass)
 	return 0;
 }
 
-/* Helper for PyObject_Dir.
-   If obj has an attr named attrname that's a list, merge its string
-   elements into keys of dict.
-   Return 0 on success, -1 on error.  Errors due to not finding the attr,
-   or the attr not being a list, are suppressed.
-*/
+/* Helper for PyObject_Dir without arguments: returns the local scope. */
+static PyObject *
+_dir_locals()
+{
+	PyObject *locals = PyEval_GetLocals();
 
-static int
-merge_list_attr(PyObject* dict, PyObject* obj, const char *attrname)
+	if (locals == NULL) {
+		PyErr_SetString(PyExc_SystemError, "frame does not exist");
+		return NULL;
+	}
+
+	/* the locals don't need to be DECREF'd */
+	return PyMapping_Keys(locals);
+}
+
+/* Helper for PyObject_Dir of type objects: returns __dict__ and __bases__.
+   We deliberately don't suck up its __class__, as methods belonging to the 
+   metaclass would probably be more confusing than helpful. 
+*/
+static PyObject * 
+_specialized_dir_type(PyObject *obj)
 {
-	PyObject *list;
-	int result = 0;
+	PyObject *result = NULL;
+	PyObject *dict = PyDict_New();
 
-	assert(PyDict_Check(dict));
-	assert(obj);
-	assert(attrname);
+	if (dict != NULL && merge_class_dict(dict, obj) == 0)
+		result = PyDict_Keys(dict);
 
-	list = PyObject_GetAttrString(obj, attrname);
-	if (list == NULL)
-		PyErr_Clear();
+	Py_XDECREF(dict);
+	return result;
+}
 
-	else if (PyList_Check(list)) {
-		int i;
-		for (i = 0; i < PyList_GET_SIZE(list); ++i) {
-			PyObject *item = PyList_GET_ITEM(list, i);
-			if (PyString_Check(item)) {
-				result = PyDict_SetItem(dict, item, Py_None);
-				if (result < 0)
-					break;
-			}
+/* Helper for PyObject_Dir of module objects: returns the module's __dict__. */
+static PyObject *
+_specialized_dir_module(PyObject *obj)
+{
+	PyObject *result = NULL;
+	PyObject *dict = PyObject_GetAttrString(obj, "__dict__");
+
+	if (dict != NULL) {
+		if (PyDict_Check(dict))
+			result = PyDict_Keys(dict);
+		else {
+			PyErr_Format(PyExc_TypeError,
+				     "%.200s.__dict__ is not a dictionary",
+				     PyModule_GetName(obj));
 		}
 	}
 
-	Py_XDECREF(list);
+	Py_XDECREF(dict);
 	return result;
 }
 
-/* Like __builtin__.dir(arg).  See bltinmodule.c's builtin_dir for the
-   docstring, which should be kept in synch with this implementation. */
-
-PyObject *
-PyObject_Dir(PyObject *arg)
+/* Helper for PyObject_Dir of generic objects: returns __dict__, __class__,
+   and recursively up the __class__.__bases__ chain.
+*/
+static PyObject *
+_generic_dir(PyObject *obj)
 {
-	/* Set exactly one of these non-NULL before the end. */
-	PyObject *result = NULL;	/* result list */
-	PyObject *masterdict = NULL;	/* result is masterdict.keys() */
-
-	/* If NULL arg, return the locals. */
-	if (arg == NULL) {
-		PyObject *locals = PyEval_GetLocals();
-		if (locals == NULL)
-			goto error;
-		result = PyMapping_Keys(locals);
-		if (result == NULL)
-			goto error;
+	PyObject *result = NULL;
+	PyObject *dict = NULL;
+	PyObject *itsclass = NULL;
+	
+	/* Get __dict__ (which may or may not be a real dict...) */
+	dict = PyObject_GetAttrString(obj, "__dict__");
+	if (dict == NULL) {
+		PyErr_Clear();
+		dict = PyDict_New();
 	}
-
-	/* Elif this is some form of module, we only want its dict. */
-	else if (PyModule_Check(arg)) {
-		masterdict = PyObject_GetAttrString(arg, "__dict__");
-		if (masterdict == NULL)
-			goto error;
-		if (!PyDict_Check(masterdict)) {
-			PyErr_SetString(PyExc_TypeError,
-					"module.__dict__ is not a dictionary");
-			goto error;
-		}
+	else if (!PyDict_Check(dict)) {
+		Py_DECREF(dict);
+		dict = PyDict_New();
 	}
-
-	/* Elif some form of type or class, grab its dict and its bases.
-	   We deliberately don't suck up its __class__, as methods belonging
-	   to the metaclass would probably be more confusing than helpful. */
-	else if (PyType_Check(arg)) {
-		masterdict = PyDict_New();
-		if (masterdict == NULL)
-			goto error;
-		if (merge_class_dict(masterdict, arg) < 0)
-			goto error;
+	else {
+		/* Copy __dict__ to avoid mutating it. */
+		PyObject *temp = PyDict_Copy(dict);
+		Py_DECREF(dict);
+		dict = temp;
 	}
 
-	/* Else look at its dict, and the attrs reachable from its class. */
+	if (dict == NULL)
+		goto error;
+
+	/* Merge in attrs reachable from its class. */
+	itsclass = PyObject_GetAttrString(obj, "__class__");
+	if (itsclass == NULL)
+		/* XXX(tomer): Perhaps fall back to obj->ob_type if no
+		               __class__ exists? */
+		PyErr_Clear();
 	else {
-		PyObject *itsclass;
-		/* Create a dict to start with.  CAUTION:  Not everything
-		   responding to __dict__ returns a dict! */
-		masterdict = PyObject_GetAttrString(arg, "__dict__");
-		if (masterdict == NULL) {
-			PyErr_Clear();
-			masterdict = PyDict_New();
-		}
-		else if (!PyDict_Check(masterdict)) {
-			Py_DECREF(masterdict);
-			masterdict = PyDict_New();
-		}
-		else {
-			/* The object may have returned a reference to its
-			   dict, so copy it to avoid mutating it. */
-			PyObject *temp = PyDict_Copy(masterdict);
-			Py_DECREF(masterdict);
-			masterdict = temp;
-		}
-		if (masterdict == NULL)
+		if (merge_class_dict(dict, itsclass) != 0)
 			goto error;
+	}
 
-		/* Merge in __members__ and __methods__ (if any).
-		   XXX Would like this to go away someday; for now, it's
-		   XXX needed to get at im_self etc of method objects. */
-		if (merge_list_attr(masterdict, arg, "__members__") < 0)
-			goto error;
-		if (merge_list_attr(masterdict, arg, "__methods__") < 0)
-			goto error;
+	result = PyDict_Keys(dict);
+	/* fall through */
+error:
+	Py_XDECREF(itsclass);
+	Py_XDECREF(dict);
+	return result;
+}
 
-		/* Merge in attrs reachable from its class.
-		   CAUTION:  Not all objects have a __class__ attr. */
-		itsclass = PyObject_GetAttrString(arg, "__class__");
-		if (itsclass == NULL)
-			PyErr_Clear();
-		else {
-			int status = merge_class_dict(masterdict, itsclass);
-			Py_DECREF(itsclass);
-			if (status < 0)
-				goto error;
-		}
-	}
+/* Helper for PyObject_Dir: object introspection.
+   This calls one of the above specialized versions if no __dir__ method
+   exists. */
+static PyObject *
+_dir_object(PyObject *obj)
+{
+	PyObject * result = NULL;
+	PyObject * dirfunc = PyObject_GetAttrString((PyObject*)obj->ob_type,
+						    "__dir__");
 
-	assert((result == NULL) ^ (masterdict == NULL));
-	if (masterdict != NULL) {
-		/* The result comes from its keys. */
-		assert(result == NULL);
-		result = PyDict_Keys(masterdict);
-		if (result == NULL)
-			goto error;
+	assert(obj);
+	if (dirfunc == NULL) {
+		/* use default implementation */
+		PyErr_Clear();
+		if (PyModule_Check(obj))
+			result = _specialized_dir_module(obj);
+		else if (PyType_Check(obj))
+			result = _specialized_dir_type(obj);
+		else
+			result = _generic_dir(obj);
 	}
+	else {
+		/* use __dir__ */
+		result = PyObject_CallFunctionObjArgs(dirfunc, obj, NULL);
+		Py_DECREF(dirfunc);
+		if (result == NULL)
+			return NULL;
 
-	assert(result);
-	if (!PyList_Check(result)) {
-		PyErr_Format(PyExc_TypeError,
-			"Expected keys() to be a list, not '%.200s'",
-			result->ob_type->tp_name);
-		goto error;
+		/* result must be a list */
+		/* XXX(gbrandl): could also check if all items are strings */
+		if (!PyList_Check(result)) {
+			PyErr_Format(PyExc_TypeError,
+				     "__dir__() must return a list, not %.200s",
+				     result->ob_type->tp_name);
+			Py_DECREF(result);
+			result = NULL;
+		}
 	}
-	if (PyList_Sort(result) != 0)
-		goto error;
+
+	return result;
+}
+
+/* Implementation of dir() -- if obj is NULL, returns the names in the current
+   (local) scope.  Otherwise, performs introspection of the object: returns a
+   sorted list of attribute names (supposedly) accessible from the object
+*/
+PyObject *
+PyObject_Dir(PyObject *obj)
+{
+	PyObject * result;
+
+	if (obj == NULL)
+		/* no object -- introspect the locals */
+		result = _dir_locals();
 	else
-		goto normal_return;
+		/* object -- introspect the object */
+		result = _dir_object(obj);
 
-  error:
-	Py_XDECREF(result);
-	result = NULL;
-	/* fall through */
-  normal_return:
-  	Py_XDECREF(masterdict);
+	assert(result == NULL || PyList_Check(result));
+
+	if (result != NULL && PyList_Sort(result) != 0) {
+		/* sorting the list failed */
+		Py_DECREF(result);
+		result = NULL;
+	}
+	
 	return result;
 }