builtin: Implement builtin_iter

corona10 · corona10 · commit 29323cf040b6 · 2019-02-14T02:53:02.000+09:00
diff --git a/builtin/builtin.go b/builtin/builtin.go
@@ -44,7 +44,7 @@ func init() {
 		// py.MustNewMethod("input", builtin_input, 0, input_doc),
 		// py.MustNewMethod("isinstance", builtin_isinstance, 0, isinstance_doc),
 		// py.MustNewMethod("issubclass", builtin_issubclass, 0, issubclass_doc),
-		// py.MustNewMethod("iter", builtin_iter, 0, iter_doc),
+		py.MustNewMethod("iter", builtin_iter, 0, iter_doc),
 		py.MustNewMethod("len", builtin_len, 0, len_doc),
 		py.MustNewMethod("locals", py.InternalMethodLocals, 0, locals_doc),
 		py.MustNewMethod("max", builtin_max, 0, max_doc),
@@ -762,6 +762,39 @@ object.
 The globals and locals are dictionaries, defaulting to the current
 globals and locals.  If only globals is given, locals defaults to it.`
 
+const iter_doc = `iter(iterable) -> iterator
+iter(callable, sentinel) -> iterator
+
+Get an iterator from an object.  In the first form, the argument must
+supply its own iterator, or be a sequence.
+In the second form, the callable is called until it returns the sentinel.
+`
+
+func builtin_iter(self py.Object, args py.Tuple) (py.Object, error) {
+	nArgs := len(args)
+	if nArgs < 1 {
+		return nil, py.ExceptionNewf(py.TypeError,
+			"iter expected at least 1 arguments, got %d",
+			nArgs)
+	} else if nArgs > 2 {
+		return nil, py.ExceptionNewf(py.TypeError,
+			"iter expected at most 2 arguments, got %d",
+			nArgs)
+	}
+
+	v := args[0]
+	if nArgs == 1 {
+		return py.Iter(v)
+	}
+	_, ok := v.(*py.Function)
+	sentinel := args[1]
+	if !ok {
+		return nil, py.ExceptionNewf(py.TypeError,
+			"iter(v, w): v must be callable")
+	}
+	return py.NewCallIterator(v, sentinel), nil
+}
+
 // For code see vm/builtin.go
 
 const len_doc = `len(object) -> integer
diff --git a/builtin/tests/builtin.py b/builtin/tests/builtin.py
@@ -132,6 +132,50 @@ def func(p):
     ok = True
 assert ok, "ValueError not raised"
 
+doc="iter"
+cnt = 0
+def f():
+    global cnt
+    cnt += 1
+    return cnt
+
+l = list(iter(f,20))
+assert len(l) == 19
+for idx, v in enumerate(l):
+    assert idx + 1 == v
+
+words1 = ['g', 'p', 'y', 't', 'h', 'o', 'n']
+words2 = list(iter(words1))
+for w1, w2 in zip(words1, words2):
+    assert w1 == w2
+
+ok = False
+try:
+    iter()
+except TypeError:
+    ok = True
+finally:
+    assert ok, "TypeError not raised"
+    ok = False
+
+try:
+    l = [1, 2, 3]
+    iter(l, 2)
+except TypeError:
+    ok = True
+finally:
+    assert ok, "TypeError not raised"
+    ok = False
+
+try:
+    iter(f, 2, 3)
+except TypeError:
+    ok = True
+finally:
+    assert ok, "TypeError not raised"
+    ok = False
+
+
 doc="next no default"
 def gen():
     yield 1
diff --git a/py/call_iterator.go b/py/call_iterator.go
@@ -0,0 +1,55 @@
+// Copyright 2019 The go-python Authors.  All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// CallIterator objects
+
+package py
+
+// A python CallIterator object
+type CallIterator struct {
+	callable Object
+	sentinel Object
+}
+
+var CallIteratorType = NewType("callable_iterator", "callable_iterator type")
+
+// Type of this object
+func (o *CallIterator) Type() *Type {
+	return CallIteratorType
+}
+
+func (cit *CallIterator) M__iter__() (Object, error) {
+	return cit, nil
+}
+
+// Get next one from the iteration
+func (cit *CallIterator) M__next__() (Object, error) {
+	f, ok := cit.callable.(*Function)
+	if !ok {
+		return nil, ExceptionNewf(TypeError, "'%s' object is not callable", f.Type())
+	}
+	value, err := Call(f, nil, nil)
+
+	if err != nil {
+		return nil, err
+	}
+
+	if value == cit.sentinel {
+		return nil, StopIteration
+	}
+
+	return value, nil
+}
+
+// Define a new CallIterator
+func NewCallIterator(callable Object, sentinel Object) *CallIterator {
+	c := &CallIterator{
+		callable: callable,
+		sentinel: sentinel,
+	}
+	return c
+}
+
+// Check interface is satisfied
+var _ I_iterator = (*CallIterator)(nil)
diff --git a/py/tests/iter.py b/py/tests/iter.py
@@ -0,0 +1,21 @@
+# Copyright 2019 The go-python Authors.  All rights reserved.
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file.
+
+doc="iter"
+cnt = 0
+def f():
+    global cnt
+    cnt += 1
+    return cnt
+
+l = list(iter(f,20))
+assert len(l) == 19
+for idx, v in enumerate(l):
+    assert idx + 1 == v
+
+words1 = ['g', 'p', 'y', 't', 'h', 'o', 'n']
+words2 = list(iter(words1))
+for w1, w2 in zip(words1, words2):
+    assert w1 == w2
+doc="finished"
