Author: Armin Rigo <ar...@tunes.org>
Branch: stmgc-c8-dictiter
Changeset: r80557:cf95b2a3e1ca
Date: 2015-11-05 22:59 +0100
http://bitbucket.org/pypy/pypy/changeset/cf95b2a3e1ca/
Log: stmset.__iter__
diff --git a/pypy/module/pypystm/stmset.py b/pypy/module/pypystm/stmset.py
--- a/pypy/module/pypystm/stmset.py
+++ b/pypy/module/pypystm/stmset.py
@@ -2,6 +2,7 @@
The class pypystm.stmset, giving a part of the regular 'set' interface
"""
+from pypy.interpreter.error import OperationError
from pypy.interpreter.baseobjspace import W_Root
from pypy.interpreter.typedef import TypeDef
from pypy.interpreter.gateway import interp2app
@@ -150,8 +151,48 @@
return space.wrap(self.get_length())
def iter_w(self, space):
- # not a real lazy iterator!
- return space.iter(space.newlist(self.get_items_w()))
+ return W_STMSetIter(self.h)
+
+
+class W_STMSetIter(W_Root):
+ _immutable_fields_ = ["hiter"]
+ next_from_same_hash = 0
+
+ def __init__(self, hobj):
+ self.hiter = hobj.iterentries()
+
+ def descr_iter(self, space):
+ return self
+
+ def descr_length_hint(self, space):
+ # xxx estimate: doesn't remove the items already yielded,
+ # and uses the faster len_estimate(); on the other hand,
+ # counts only one for every 64-bit hash value
+ return space.wrap(self.hiter.hashtable.len_estimate())
+
+ def descr_next(self, space):
+ if self.next_from_same_hash == 0: # common case
+ try:
+ entry = self.hiter.next()
+ except StopIteration:
+ raise OperationError(space.w_StopIteration, space.w_None)
+ index = 0
+ array = lltype.cast_opaque_ptr(PARRAY, entry.object)
+ else:
+ index = self.next_from_same_hash
+ array = self.next_array
+ self.next_from_same_hash = 0
+ self.next_array = lltype.nullptr(ARRAY)
+ #
+ if len(array) > index + 1: # uncommon case
+ self.next_from_same_hash = index + 1
+ self.next_array = array
+ #
+ return cast_gcref_to_instance(W_Root, array[index])
+
+ def _cleanup_(self):
+ raise Exception("seeing a prebuilt %r object" % (
+ self.__class__,))
def W_STMSet___new__(space, w_subtype):
@@ -170,3 +211,10 @@
__len__ = interp2app(W_STMSet.len_w),
__iter__ = interp2app(W_STMSet.iter_w),
)
+
+W_STMSetIter.typedef = TypeDef(
+ "stmset_iter",
+ __iter__ = interp2app(W_STMSetIter.descr_iter),
+ next = interp2app(W_STMSetIter.descr_next),
+ __length_hint__ = interp2app(W_STMSetIter.descr_length_hint),
+ )
diff --git a/pypy/module/pypystm/test/test_stmset.py
b/pypy/module/pypystm/test/test_stmset.py
--- a/pypy/module/pypystm/test/test_stmset.py
+++ b/pypy/module/pypystm/test/test_stmset.py
@@ -83,3 +83,19 @@
assert len(s) == 2
items = list(s)
assert items == [42.5, key3] or items == [key3, 42.5]
+
+ def test_iterator(self):
+ import pypystm
+ class A(object):
+ def __hash__(self):
+ return 42
+ class B(object):
+ pass
+ d = pypystm.stmset()
+ a1 = A()
+ a2 = A()
+ b0 = B()
+ d.add(a1)
+ d.add(a2)
+ d.add(b0)
+ assert sorted(d) == sorted([a1, a2, b0])
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit
