py: Support zip builtin feature

corona10 · ncw · commit d8a0825f736e · 2019-01-03T18:04:36.000Z
diff --git a/builtin/builtin.go b/builtin/builtin.go
@@ -92,7 +92,7 @@ func init() {
 		// "super":          py.SuperType,
 		"tuple": py.TupleType,
 		"type":  py.TypeType,
-		// "zip":            py.ZipType,
+		"zip":   py.ZipType,
 
 		// Exceptions
 		"ArithmeticError":           py.ArithmeticError,
diff --git a/builtin/tests/builtin.py b/builtin/tests/builtin.py
@@ -232,6 +232,20 @@ class C: pass
     ok = True
 assert ok, "TypeError not raised"
 
+doc="zip"
+ok = False
+a = [3, 4, 5, 6, 7]
+b = [8, 9, 10, 11, 12]
+assert [e for e in zip(a, b)] == [(3,8), (4,9), (5,10), (6,11), (7,12)]
+try:
+    zip(1,2,3)
+except TypeError as e:
+    print(e.args[0])
+    if e.args[0] != "zip argument #1 must support iteration":
+        raise
+    ok = True
+assert ok, "TypeError not raised"
+
 doc="__import__"
 lib = __import__("lib")
 assert lib.libfn() == 42
diff --git a/py/tests/zip.py b/py/tests/zip.py
@@ -0,0 +1,13 @@
+# 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.
+
+doc="zip"
+a = [1, 2, 3, 4, 5]
+b = [10, 11, 12]
+c = [e for e in zip(a, b)]
+assert len(c) == 3
+for idx, e in enumerate(c):
+    assert a[idx] == c[idx][0]
+    assert b[idx] == c[idx][1]
+doc="finished"
diff --git a/py/zip.go b/py/zip.go
@@ -0,0 +1,66 @@
+// 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 Zip object
+type Zip struct {
+	itTuple Tuple
+	size    int
+}
+
+// A python ZipIterator iterator
+type ZipIterator struct {
+	zip Zip
+}
+
+var ZipType = NewTypeX("zip", `zip(iter1 [,iter2 [...]]) --> zip object
+
+Return a zip object whose .__next__() method returns a tuple where
+the i-th element comes from the i-th iterable argument.  The .__next__()
+method continues until the shortest iterable in the argument sequence
+is exhausted and then it raises StopIteration.`,
+	ZipTypeNew, nil)
+
+// Type of this object
+func (z *Zip) Type() *Type {
+	return ZipType
+}
+
+// ZipTypeNew
+func ZipTypeNew(metatype *Type, args Tuple, kwargs StringDict) (Object, error) {
+	tupleSize := len(args)
+	itTuple := make(Tuple, tupleSize)
+	for i := 0; i < tupleSize; i++ {
+		item := args[i]
+		iter, err := Iter(item)
+		if err != nil {
+			return nil, ExceptionNewf(TypeError, "zip argument #%d must support iteration", i+1)
+		}
+		itTuple[i] = iter
+	}
+
+	return &Zip{itTuple: itTuple, size: tupleSize}, nil
+}
+
+// Zip iterator
+func (z *Zip) M__iter__() (Object, error) {
+	return z, nil
+}
+
+func (z *Zip) M__next__() (Object, error) {
+	result := make(Tuple, z.size)
+	for i := 0; i < z.size; i++ {
+		value, err := Next(z.itTuple[i])
+		if err != nil {
+			return nil, err
+		}
+		result[i] = value
+	}
+	return result, nil
+}
+
+// Check interface is satisfied
+var _ I__iter__ = (*Zip)(nil)
+var _ I__next__ = (*Zip)(nil)
