[pypy-commit] pypy iterator-in-rpython: more support

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: iterator-in-rpython
Changeset: r56002:1fb746e47747
Date: 2012-07-08 19:21 +0200
http://bitbucket.org/pypy/pypy/changeset/1fb746e47747/

Log:	more support

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
@@ -3806,6 +3806,25 @@
         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
@@ -664,6 +664,10 @@
         s_iterable = ins._true_getattr('__iter__')
         return s_iterable.call(getbookkeeper().build_args("simple_call", []))
 
+    def next(ins):
+        s_next = ins._true_getattr('next')
+        return s_next.call(getbookkeeper().build_args('simple_call', []))
+
 class __extend__(SomeBuiltin):
     def _can_only_throw(bltn, *args):
         analyser_func = getattr(bltn.analyser, 'im_func', None)
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,17 @@
     def rtype_is_true(self, hop):
         raise NotImplementedError
 
+    def rtype_iter(self, hop):
+        vinst, = hop.inputargs(self)
+        vcls = self.getfield(vinst, '__class__', hop.llops)
+        if '__iter__' not in self.rclass.allmethods:
+            raise Exception("Only supporting iterators with __iter__ as a method")
+        viter = self.rclass.getclsfield(vcls, '__iter__', hop.llops)
+        return hop.gendirectcall(viter, vinst)
+
+    def rtype_next(self, hop):
+        xxx
+
     def ll_str(self, i):
         raise NotImplementedError