[pypy-commit] pypy unicode-strategies: implement iterator for unicode sets

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: unicode-strategies
Changeset: r58477:30b58451e900
Date: 2012-10-26 21:23 +0200
http://bitbucket.org/pypy/pypy/changeset/30b58451e900/

Log:	implement iterator for unicode sets

diff --git a/pypy/objspace/std/setobject.py b/pypy/objspace/std/setobject.py
--- a/pypy/objspace/std/setobject.py
+++ b/pypy/objspace/std/setobject.py
@@ -748,8 +748,7 @@
         return self.space.wrap(item)
 
     def iter(self, w_set):
-        XXX
-        return StringIteratorImplementation(self.space, self, w_set)
+        return UnicodeIteratorImplementation(self.space, self, w_set)
 
 
 class IntegerSetStrategy(AbstractUnwrappedSetStrategy, SetStrategy):
@@ -891,6 +890,18 @@
         else:
             return None
 
+class UnicodeIteratorImplementation(IteratorImplementation):
+    def __init__(self, space, strategy, w_set):
+        IteratorImplementation.__init__(self, space, strategy, w_set)
+        d = strategy.unerase(w_set.sstorage)
+        self.iterator = d.iterkeys()
+
+    def next_entry(self):
+        for key in self.iterator:
+            return self.space.wrap(key)
+        else:
+            return None
+
 class IntegerIteratorImplementation(IteratorImplementation):
     #XXX same implementation in dictmultiobject on dictstrategy-branch
     def __init__(self, space, strategy, w_set):
diff --git a/pypy/objspace/std/test/test_setstrategies.py b/pypy/objspace/std/test/test_setstrategies.py
--- a/pypy/objspace/std/test/test_setstrategies.py
+++ b/pypy/objspace/std/test/test_setstrategies.py
@@ -1,7 +1,10 @@
 from pypy.objspace.std.setobject import W_SetObject
 from pypy.objspace.std.setobject import (IntegerSetStrategy, ObjectSetStrategy,
                                          EmptySetStrategy, StringSetStrategy,
-                                         UnicodeSetStrategy)
+                                         UnicodeSetStrategy,
+                                         IntegerIteratorImplementation,
+                                         StringIteratorImplementation,
+                                         UnicodeIteratorImplementation)
 from pypy.objspace.std.listobject import W_ListObject
 
 class TestW_SetStrategies:
@@ -118,3 +121,23 @@
 
         assert s1.has_key(self.space.wrap(FakeInt(2)))
         assert s1.strategy is self.space.fromcache(ObjectSetStrategy)
+
+    def test_iter(self):
+        space = self.space
+        s = W_SetObject(space, self.wrapped([1,2]))
+        it = s.iter()
+        assert isinstance(it, IntegerIteratorImplementation)
+        assert space.unwrap(it.next()) == 1
+        assert space.unwrap(it.next()) == 2
+        #
+        s = W_SetObject(space, self.wrapped(["a", "b"]))
+        it = s.iter()
+        assert isinstance(it, StringIteratorImplementation)
+        assert space.unwrap(it.next()) == "a"
+        assert space.unwrap(it.next()) == "b"
+        #
+        s = W_SetObject(space, self.wrapped([u"a", u"b"]))
+        it = s.iter()
+        assert isinstance(it, UnicodeIteratorImplementation)
+        assert space.unwrap(it.next()) == u"a"
+        assert space.unwrap(it.next()) == u"b"