builtin: Implement enumerate feature

corona10 · corona10 · commit cfd27cb7ccfd · 2018-12-25T20:31:12.000+09:00
Now, gpython supports enumerate feature
diff --git a/builtin/builtin.go b/builtin/builtin.go
@@ -74,7 +74,7 @@ func init() {
 		"classmethod": py.ClassMethodType,
 		"complex":     py.ComplexType,
 		"dict":        py.StringDictType, // FIXME
-		// "enumerate":      py.EnumType,
+		"enumerate":   py.EnumerateType,
 		// "filter":         py.FilterType,
 		"float":     py.FloatType,
 		"frozenset": py.FrozenSetType,
diff --git a/builtin/tests/builtin.py b/builtin/tests/builtin.py
@@ -32,6 +32,14 @@
 doc="divmod"
 assert divmod(34,7) == (4, 6)
 
+doc="enumerate"
+a = [e for e in enumerate([3,4,5,6,7], 4)]
+idxs = [4, 5, 6, 7, 8]
+values = [3, 4, 5, 6, 7]
+for idx, value in enumerate(values):
+    assert idxs[idx] == a[idx][0]
+    assert values[idx] == a[idx][1]
+
 doc="eval"
 # smoke test only - see vm/tests/builtin.py for more tests
 assert eval("1+2") == 3
diff --git a/py/enumerate.go b/py/enumerate.go
@@ -0,0 +1,88 @@
+// Copyright 2018 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.
+
+package py
+
+// A python Enumerate object
+type Enumerate struct {
+	Iterable Object
+	Start    Int
+}
+
+// A python Enumerate iterator
+type EnumerateIterator struct {
+	Enumerate
+	Index Int
+}
+
+var EnumerateType = NewTypeX("enumerate", `enumerate(iterable, start=0)
+
+Return an enumerate object.`,
+	EnumerateNew, nil)
+
+var EnumerateIteratorType = NewType("enumerate_iterator", `enumerate_iterator object`)
+
+// Type of this object
+func (o *Enumerate) Type() *Type {
+	return EnumerateType
+}
+
+// Type of this object
+func (o *EnumerateIterator) Type() *Type {
+	return EnumerateIteratorType
+}
+
+// EnumerateTypeNew
+func EnumerateNew(metatype *Type, args Tuple, kwargs StringDict) (Object, error) {
+	var iterable Object
+	var start Object
+	err := UnpackTuple(args, kwargs, "enumerate", 1, 2, &iterable, &start)
+	if err != nil {
+		return nil, err
+	}
+
+	if start == nil {
+		start = Int(0)
+	}
+	startIndex, err := Index(start)
+	if err != nil {
+		return nil, err
+	}
+	iter, err := Iter(iterable)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Enumerate{Iterable: iter, Start: startIndex}, nil
+}
+
+// Enumerate iterator
+func (it *Enumerate) M__iter__() (Object, error) {
+	return &EnumerateIterator{
+		Enumerate: *it,
+		Index:     it.Start,
+	}, nil
+}
+
+// EnumerateIterator iterator
+func (it *EnumerateIterator) M__iter__() (Object, error) {
+	return it, nil
+}
+
+// EnumerateIterator iterator next
+func (iter *EnumerateIterator) M__next__() (Object, error) {
+	value, err := Next(iter.Enumerate.Iterable)
+	if err != nil {
+		return nil, err
+	}
+	res := make(Tuple, 2)
+	res[0] = iter.Index
+	res[1] = value
+	iter.Index += 1
+	return res, nil
+}
+
+// Check interface is satisfied
+var _ I__iter__ = (*Enumerate)(nil)
+var _ I_iterator = (*EnumerateIterator)(nil)
diff --git a/py/tests/list.py b/py/tests/list.py
@@ -14,4 +14,12 @@
 assert repr([1,[2,3],4]) == "[1, [2, 3], 4]"
 assert repr(["1",[2.5,17,[]]]) == "['1', [2.5, 17, []]]"
 
+doc="enumerate"
+a = [e for e in enumerate([3,4,5,6,7], 4)]
+idxs = [4, 5, 6, 7, 8]
+values = [3, 4, 5, 6, 7]
+for idx, value in enumerate(values):
+    assert idxs[idx] == a[idx][0]
+    assert values[idx] == a[idx][1]
+
 doc="finished"
