[pypy-svn] pypy collections-module: A helper: space.compare_by_iteration().

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: collections-module
Changeset: r41966:5fbfabfab60a
Date: 2011-02-15 14:49 +0100
http://bitbucket.org/pypy/pypy/changeset/5fbfabfab60a/

Log:	A helper: space.compare_by_iteration().

diff --git a/pypy/interpreter/baseobjspace.py b/pypy/interpreter/baseobjspace.py
--- a/pypy/interpreter/baseobjspace.py
+++ b/pypy/interpreter/baseobjspace.py
@@ -1009,6 +1009,38 @@
         return self.call_args(w_func, args)
     appexec._annspecialcase_ = 'specialize:arg(2)'
 
+    def _next_or_none(self, w_it):
+        try:
+            return self.next(w_it)
+        except OperationError, e:
+            if not e.match(self, self.w_StopIteration):
+                raise
+            return None
+
+    def compare_by_iteration(self, w_iterable1, w_iterable2, op):
+        w_it1 = self.iter(w_iterable1)
+        w_it2 = self.iter(w_iterable2)
+        while True:
+            w_x1 = self._next_or_none(w_it1)
+            w_x2 = self._next_or_none(w_it2)
+            if w_x1 is None or w_x2 is None:
+                if op == 'eq': return self.newbool(w_x1 is w_x2)  # both None
+                if op == 'ne': return self.newbool(w_x1 is not w_x2)
+                if op == 'lt': return self.newbool(w_x2 is not None)
+                if op == 'le': return self.newbool(w_x1 is None)
+                if op == 'gt': return self.newbool(w_x1 is not None)
+                if op == 'ge': return self.newbool(w_x2 is None)
+                assert False, "bad value for op"
+            if not self.eq_w(w_x1, w_x2):
+                if op == 'eq': return self.w_False
+                if op == 'ne': return self.w_True
+                if op == 'lt': return self.lt(w_x1, w_x2)
+                if op == 'le': return self.le(w_x1, w_x2)
+                if op == 'gt': return self.gt(w_x1, w_x2)
+                if op == 'ge': return self.ge(w_x1, w_x2)
+                assert False, "bad value for op"
+    compare_by_iteration._annspecialcase_ = 'specialize:arg(3)'
+
     def decode_index(self, w_index_or_slice, seqlength):
         """Helper for custom sequence implementations
              -> (index, 0, 0) or

diff --git a/pypy/interpreter/test/test_objspace.py b/pypy/interpreter/test/test_objspace.py
--- a/pypy/interpreter/test/test_objspace.py
+++ b/pypy/interpreter/test/test_objspace.py
@@ -243,6 +243,21 @@
         w_res = space.call_obj_args(w_a, w_9, Arguments(space, []))        
         assert w_res is w_9
 
+    def test_compare_by_iteration(self):
+        import operator
+        space = self.space
+        for op in ['eq', 'ne', 'lt', 'le', 'gt', 'ge']:
+            comparer = getattr(operator, op)
+            for x in [[], [0], [0, 1], [0, 1, 2]]:
+                for y in [[], [-1], [0], [1], [-1, 0],
+                          [0, 0], [0, 1], [0, 2],
+                          [0, 1, 1], [0, 1, 2], [0, 1, 3]]:
+                        w_res = space.compare_by_iteration(space.wrap(x),
+                                                           space.wrap(y), op)
+                        if comparer(x, y):
+                            assert w_res is space.w_True
+                        else:
+                            assert w_res is space.w_False
 
 class TestModuleMinimal: 
     def test_sys_exists(self):