[pypy-commit] pypy default: Merge iterator-in-rpython. Now you can use __iter__ in rpython!

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r56030:ea1bfc0445f5
Date: 2012-07-11 18:39 +0200
http://bitbucket.org/pypy/pypy/changeset/ea1bfc0445f5/

Log:	Merge iterator-in-rpython. Now you can use __iter__ in rpython!

diff --git a/pypy/annotation/bookkeeper.py b/pypy/annotation/bookkeeper.py
--- a/pypy/annotation/bookkeeper.py
+++ b/pypy/annotation/bookkeeper.py
@@ -201,6 +201,7 @@
                     for op in block.operations:
                         if op.opname in ('simple_call', 'call_args'):
                             yield op
+
                         # some blocks are partially annotated
                         if binding(op.result, None) is None:
                             break   # ignore the unannotated part
diff --git a/pypy/annotation/test/test_annrpython.py b/pypy/annotation/test/test_annrpython.py
--- a/pypy/annotation/test/test_annrpython.py
+++ b/pypy/annotation/test/test_annrpython.py
@@ -3793,7 +3793,37 @@
         assert isinstance(s, annmodel.SomeString)
         assert s.no_nul
 
-
+    def test_base_iter(self):
+        class A(object):
+            def __iter__(self):
+                return self
+        
+        def fn():
+            return iter(A())
+
+        a = self.RPythonAnnotator()
+        s = a.build_types(fn, [])
+        assert isinstance(s, annmodel.SomeInstance)
+        assert s.classdef.name.endswith('.A')
+
+    def test_iter_next(self):
+        class A(object):
+            def __iter__(self):
+                return self
+
+            def next(self):
+                return 1
+        
+        def fn():
+            s = 0
+            for x in A():
+                s += x
+            return s
+
+        a = self.RPythonAnnotator()
+        s = a.build_types(fn, [])
+        assert len(a.translator.graphs) == 3 # fn, __iter__, next
+        assert isinstance(s, annmodel.SomeInteger)
 
 def g(n):
     return [0,1,2,n]
diff --git a/pypy/annotation/unaryop.py b/pypy/annotation/unaryop.py
--- a/pypy/annotation/unaryop.py
+++ b/pypy/annotation/unaryop.py
@@ -609,33 +609,36 @@
 
 class __extend__(SomeInstance):
 
+    def _true_getattr(ins, attr):
+        if attr == '__class__':
+            return ins.classdef.read_attr__class__()
+        attrdef = ins.classdef.find_attribute(attr)
+        position = getbookkeeper().position_key
+        attrdef.read_locations[position] = True
+        s_result = attrdef.getvalue()
+        # hack: if s_result is a set of methods, discard the ones
+        #       that can't possibly apply to an instance of ins.classdef.
+        # XXX do it more nicely
+        if isinstance(s_result, SomePBC):
+            s_result = ins.classdef.lookup_filter(s_result, attr,
+                                                  ins.flags)
+        elif isinstance(s_result, SomeImpossibleValue):
+            ins.classdef.check_missing_attribute_update(attr)
+            # blocking is harmless if the attribute is explicitly listed
+            # in the class or a parent class.
+            for basedef in ins.classdef.getmro():
+                if basedef.classdesc.all_enforced_attrs is not None:
+                    if attr in basedef.classdesc.all_enforced_attrs:
+                        raise HarmlesslyBlocked("get enforced attr")
+        elif isinstance(s_result, SomeList):
+            s_result = ins.classdef.classdesc.maybe_return_immutable_list(
+                attr, s_result)
+        return s_result
+
     def getattr(ins, s_attr):
         if s_attr.is_constant() and isinstance(s_attr.const, str):
             attr = s_attr.const
-            if attr == '__class__':
-                return ins.classdef.read_attr__class__()
-            attrdef = ins.classdef.find_attribute(attr)
-            position = getbookkeeper().position_key
-            attrdef.read_locations[position] = True
-            s_result = attrdef.getvalue()
-            # hack: if s_result is a set of methods, discard the ones
-            #       that can't possibly apply to an instance of ins.classdef.
-            # XXX do it more nicely
-            if isinstance(s_result, SomePBC):
-                s_result = ins.classdef.lookup_filter(s_result, attr,
-                                                      ins.flags)
-            elif isinstance(s_result, SomeImpossibleValue):
-                ins.classdef.check_missing_attribute_update(attr)
-                # blocking is harmless if the attribute is explicitly listed
-                # in the class or a parent class.
-                for basedef in ins.classdef.getmro():
-                    if basedef.classdesc.all_enforced_attrs is not None:
-                        if attr in basedef.classdesc.all_enforced_attrs:
-                            raise HarmlesslyBlocked("get enforced attr")
-            elif isinstance(s_result, SomeList):
-                s_result = ins.classdef.classdesc.maybe_return_immutable_list(
-                    attr, s_result)
-            return s_result
+            return ins._true_getattr(attr)
         return SomeObject()
     getattr.can_only_throw = []
 
@@ -657,6 +660,19 @@
         if not ins.can_be_None:
             s.const = True
 
+    def iter(ins):
+        s_iterable = ins._true_getattr('__iter__')
+        bk = getbookkeeper()
+        # record for calltables
+        bk.emulate_pbc_call(bk.position_key, s_iterable, [])
+        return s_iterable.call(bk.build_args("simple_call", []))
+
+    def next(ins):
+        s_next = ins._true_getattr('next')
+        bk = getbookkeeper()
+        # record for calltables
+        bk.emulate_pbc_call(bk.position_key, s_next, [])
+        return s_next.call(bk.build_args("simple_call", []))
 
 class __extend__(SomeBuiltin):
     def _can_only_throw(bltn, *args):
diff --git a/pypy/rpython/rclass.py b/pypy/rpython/rclass.py
--- a/pypy/rpython/rclass.py
+++ b/pypy/rpython/rclass.py
@@ -378,6 +378,30 @@
     def rtype_is_true(self, hop):
         raise NotImplementedError
 
+    def _emulate_call(self, hop, meth_name):
+        vinst, = hop.inputargs(self)
+        clsdef = hop.args_s[0].classdef
+        s_unbound_attr = clsdef.find_attribute(meth_name).getvalue()
+        s_attr = clsdef.lookup_filter(s_unbound_attr, meth_name,
+                                      hop.args_s[0].flags)
+        if s_attr.is_constant():
+            xxx # does that even happen?
+        if '__iter__' in self.allinstancefields:
+            raise Exception("__iter__ on instance disallowed")
+        r_method = self.rtyper.makerepr(s_attr)
+        r_method.get_method_from_instance(self, vinst, hop.llops)
+        hop2 = hop.copy()
+        hop2.spaceop.opname = 'simple_call'
+        hop2.args_r = [r_method]
+        hop2.args_s = [s_attr]
+        return hop2.dispatch()
+
+    def rtype_iter(self, hop):
+        return self._emulate_call(hop, '__iter__')
+
+    def rtype_next(self, hop):
+        return self._emulate_call(hop, 'next')
+
     def ll_str(self, i):
         raise NotImplementedError
 
diff --git a/pypy/rpython/test/test_rclass.py b/pypy/rpython/test/test_rclass.py
--- a/pypy/rpython/test/test_rclass.py
+++ b/pypy/rpython/test/test_rclass.py
@@ -1143,6 +1143,62 @@
                                       'cast_pointer': 1,
                                       'setfield': 1}
 
+    def test_iter(self):
+        class Iterable(object):
+            def __init__(self):
+                self.counter = 0
+            
+            def __iter__(self):
+                return self
+
+            def next(self):
+                if self.counter == 5:
+                    raise StopIteration
+                self.counter += 1
+                return self.counter - 1
+
+        def f():
+            i = Iterable()
+            s = 0
+            for elem in i:
+                s += elem
+            return s
+
+        assert self.interpret(f, []) == f()
+
+    def test_iter_2_kinds(self):
+        class BaseIterable(object):
+            def __init__(self):
+                self.counter = 0
+            
+            def __iter__(self):
+                return self
+
+            def next(self):
+                if self.counter >= 5:
+                    raise StopIteration
+                self.counter += self.step
+                return self.counter - 1
+        
+        class Iterable(BaseIterable):
+            step = 1
+
+        class OtherIter(BaseIterable):
+            step = 2
+
+        def f(k):
+            if k:
+                i = Iterable()
+            else:
+                i = OtherIter()
+            s = 0
+            for elem in i:
+                s += elem
+            return s
+
+        assert self.interpret(f, [True]) == f(True)
+        assert self.interpret(f, [False]) == f(False)
+
 
 class TestOOtype(BaseTestRclass, OORtypeMixin):