[pypy-svn] r58765 - in pypy/branch/2.5-merge/pypy/module/operator: . test

[jlg at codespeak.net]

Author: jlg
Date: Tue Oct  7 16:51:22 2008
New Revision: 58765

Modified:
   pypy/branch/2.5-merge/pypy/module/operator/app_operator.py
   pypy/branch/2.5-merge/pypy/module/operator/test/test_operator.py
Log:
test and implementation of support for multiple gets in attrgetter() and itemgetter()

Modified: pypy/branch/2.5-merge/pypy/module/operator/app_operator.py
==============================================================================
--- pypy/branch/2.5-merge/pypy/module/operator/app_operator.py	(original)
+++ pypy/branch/2.5-merge/pypy/module/operator/app_operator.py	Tue Oct  7 16:51:22 2008
@@ -65,18 +65,31 @@
 
 class attrgetter(object):
 
-    def __init__(self, name):
-        self.name = name
+    def __init__(self, attr, *args):
+        self.attrs = args
+        self.attr = attr
 
     def __call__(self, obj):
-        return getattr(obj, self.name)
+        result = getattr(obj, self.attr)
+
+        if self.attrs:
+            list = [result] + [getattr(obj, attr) for attr in self.attrs]
+            return tuple(list)
+
+        return result
     
 class itemgetter(object):
 
-    def __init__(self, index):
-        self.index = index
+    def __init__(self, item, *args):
+        self.items = args
+        self.item = item
 
     def __call__(self, obj):
-        return obj[self.index]
+        result = obj[self.item]
+
+        if self.items:
+            list = [result] + [obj[item] for item in self.items]
+            return tuple(list)
+
+        return result
 
-    

Modified: pypy/branch/2.5-merge/pypy/module/operator/test/test_operator.py
==============================================================================
--- pypy/branch/2.5-merge/pypy/module/operator/test/test_operator.py	(original)
+++ pypy/branch/2.5-merge/pypy/module/operator/test/test_operator.py	Tue Oct  7 16:51:22 2008
@@ -15,3 +15,23 @@
         assert a.get3("foobar") == "b"
         assert a.getx(*(a,)) == 5
         assert a.get3(obj="foobar") == "b"
+        
+
+    def test_getter_multiple_gest(self):
+        import operator
+
+        class A(object):
+            pass
+
+        a = A()
+        a.x = 'X'
+        a.y = 'Y'
+        a.z = 'Z'
+
+        assert operator.attrgetter('x','z','y')(a) == ('X', 'Z', 'Y')
+        raises(TypeError, operator.attrgetter('x', (), 'y'), a)
+
+        data = map(str, range(20))
+        assert operator.itemgetter(2,10,5)(data) == ('2', '10', '5')
+        raises(TypeError, operator.itemgetter(2, 'x', 5), data)
+