#1826: allow dotted attribute paths in operator.attrgetter.

diff --git a/Doc/library/operator.rst b/Doc/library/operator.rst
index ea4d328ffe52dbd1d7668aa7f62087d53b920637..5527dc2c8fd0cfffd03c261f22980846b16c5fb3 100644 (file)
--- a/Doc/library/operator.rst
+++ b/Doc/library/operator.rst
@@ -499,15 +499,21 @@ expect a function argument.
 
    Return a callable object that fetches *attr* from its operand. If more than one
    attribute is requested, returns a tuple of attributes. After,
-   ``f=attrgetter('name')``, the call ``f(b)`` returns ``b.name``.  After,
-   ``f=attrgetter('name', 'date')``, the call ``f(b)`` returns ``(b.name,
+   ``f = attrgetter('name')``, the call ``f(b)`` returns ``b.name``.  After,
+   ``f = attrgetter('name', 'date')``, the call ``f(b)`` returns ``(b.name,
    b.date)``.
 
+   The attribute names can also contain dots; after ``f = attrgetter('date.month')``,
+   the call ``f(b)`` returns ``b.date.month``.
+
    .. versionadded:: 2.4
 
    .. versionchanged:: 2.5
       Added support for multiple attributes.
 
+   .. versionchanged:: 2.6
+      Added support for dotted attributes.
+
 
 .. function:: itemgetter(item[, args...])
 

diff --git a/Lib/test/test_operator.py b/Lib/test/test_operator.py
index a60ceb609ec8ddc2e5503d40f4e6a5243c36ca10..3f3ea0033e5582fa5f599fb92ae4a03790863622 100644 (file)
--- a/Lib/test/test_operator.py
+++ b/Lib/test/test_operator.py
@@ -386,6 +386,26 @@ class OperatorTestCase(unittest.TestCase):
                 raise SyntaxError
         self.failUnlessRaises(SyntaxError, operator.attrgetter('foo'), C())
 
+        # recursive gets
+        a = A()
+        a.name = 'arthur'
+        a.child = A()
+        a.child.name = 'thomas'
+        f = operator.attrgetter('child.name')
+        self.assertEqual(f(a), 'thomas')
+        self.assertRaises(AttributeError, f, a.child)
+        f = operator.attrgetter('name', 'child.name')
+        self.assertEqual(f(a), ('arthur', 'thomas'))
+        f = operator.attrgetter('name', 'child.name', 'child.child.name')
+        self.assertRaises(AttributeError, f, a)
+
+        a.child.child = A()
+        a.child.child.name = 'johnson'
+        f = operator.attrgetter('child.child.name')
+        self.assertEqual(f(a), 'johnson')
+        f = operator.attrgetter('name', 'child.name', 'child.child.name')
+        self.assertEqual(f(a), ('arthur', 'thomas', 'johnson'))
+
     def test_itemgetter(self):
         a = 'ABCDE'
         f = operator.itemgetter(2)

diff --git a/Misc/NEWS b/Misc/NEWS
index 9e4a48b308b0553bb2784c9afccfe7e80a7b7612..9146bcb8e48df8944857c6c77c5556abe5c92f5c 100644 (file)
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -1196,6 +1196,8 @@ Library
 Extension Modules
 -----------------
 
+- Patch #1826: operator.attrgetter() now supports dotted attribute paths.
+
 - Patch #1957: syslogmodule: Release GIL when calling syslog(3)
 
 - #2112: mmap.error is now a subclass of EnvironmentError and not a

diff --git a/Modules/operator.c b/Modules/operator.c
index b6c3d000e8025389b319ee4e1185411c838e9bae..4d4a9401b62f5c5056b41e27b192ce53c674efca 100644 (file)
--- a/Modules/operator.c
+++ b/Modules/operator.c
@@ -495,6 +495,49 @@ attrgetter_traverse(attrgetterobject *ag, visitproc visit, void *arg)
        return 0;
 }
 
+static PyObject *
+dotted_getattr(PyObject *obj, PyObject *attr)
+{
+       char *s, *p;
+
+#ifdef Py_USING_UNICODE
+       if (PyUnicode_Check(attr)) {
+               attr = _PyUnicode_AsDefaultEncodedString(attr, NULL);
+               if (attr == NULL)
+                       return NULL;
+       }
+#endif
+       
+       if (!PyString_Check(attr)) {
+               PyErr_SetString(PyExc_TypeError,
+                               "attribute name must be a string");
+               return NULL;
+       }
+
+       s = PyString_AS_STRING(attr);
+       Py_INCREF(obj);
+       for (;;) {
+               PyObject *newobj, *str;
+               p = strchr(s, '.');
+               str = p ? PyString_FromStringAndSize(s, (p-s)) : 
+                         PyString_FromString(s);
+               if (str == NULL) {
+                       Py_DECREF(obj);
+                       return NULL;
+               }
+               newobj = PyObject_GetAttr(obj, str);
+               Py_DECREF(str);
+               Py_DECREF(obj);
+               if (newobj == NULL)
+                       return NULL;
+               obj = newobj;
+               if (p == NULL) break;
+               s = p+1;
+       }
+
+       return obj;
+}
+
 static PyObject *
 attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw)
 {
@@ -504,7 +547,7 @@ attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw)
        if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &obj))
                return NULL;
        if (ag->nattrs == 1)
-               return PyObject_GetAttr(obj, ag->attr);
+               return dotted_getattr(obj, ag->attr);
 
        assert(PyTuple_Check(ag->attr));
        assert(PyTuple_GET_SIZE(ag->attr) == nattrs);
@@ -516,7 +559,7 @@ attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw)
        for (i=0 ; i < nattrs ; i++) {
                PyObject *attr, *val;
                attr = PyTuple_GET_ITEM(ag->attr, i);
-               val = PyObject_GetAttr(obj, attr);
+               val = dotted_getattr(obj, attr);
                if (val == NULL) {
                        Py_DECREF(result);
                        return NULL;
@@ -531,7 +574,9 @@ PyDoc_STRVAR(attrgetter_doc,
 \n\
 Return a callable object that fetches the given attribute(s) from its operand.\n\
 After, f=attrgetter('name'), the call f(r) returns r.name.\n\
-After, g=attrgetter('name', 'date'), the call g(r) returns (r.name, r.date).");
+After, g=attrgetter('name', 'date'), the call g(r) returns (r.name, r.date).\n\
+After, h=attrgetter('name.first', 'name.last'), the call h(r) returns\n\
+(r.name.first, r.name.last).");
 
 static PyTypeObject attrgetter_type = {
        PyVarObject_HEAD_INIT(NULL, 0)