diff --git a/builtin/builtin.go b/builtin/builtin.go index 3486dd6e..0aea4767 100644 --- 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 index ef0d01ac..d7140d68 100644 --- 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 new file mode 100644 index 00000000..adc68d9a --- /dev/null +++ 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" \ No newline at end of file diff --git a/py/zip.go b/py/zip.go new file mode 100644 index 00000000..2626ec01 --- /dev/null +++ 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)