Implement builtin map (#97)

runtime/builtin_types.go

@@ -193,6 +193,38 @@ func builtinAbs(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
 	return Abs(f, args[0])
 }
 
+func builtinMapFn(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
+	argc := len(args)
+	if argc < 2 {
+		return nil, f.RaiseType(TypeErrorType, "map() requires at least two args")
+	}
+
+	pred := toFunctionUnsafe(args[0])
+	result := make([]*Object, 0, 2)
+
+	z, raised := zipLongest(f, args[1:])
+	if raised != nil {
+		return nil, raised
+	}
+	for _, tuple := range z {
+		if args[0] == None {
+			if argc == 2 {
+				result = append(result, tuple[0])
+			} else {
+				result = append(result, NewTuple(tuple...).ToObject())
+			}
+		} else {
+			ret, raised := pred.Call(f, tuple, nil)
+			if raised != nil {
+				return nil, raised
+			}
+			result = append(result, ret)
+		}
+	}
+
+	return NewList(result...).ToObject(), nil
+}
+
 func builtinAll(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
 	if raised := checkFunctionArgs(f, "all", args, ObjectType); raised != nil {
 		return nil, raised
@@ -571,13 +603,9 @@ func builtinZip(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
 		return NewList().ToObject(), nil
 	}
 	result := make([]*Object, 0, 2)
-	iters := make([]*Object, argc)
-	for i, arg := range args {
-		iter, raised := Iter(f, arg)
-		if raised != nil {
-			return nil, raised
-		}
-		iters[i] = iter
+	iters, raised := initIters(f, args)
+	if raised != nil {
+		return nil, raised
 	}
 
 Outer:
@@ -621,6 +649,7 @@ func init() {
 		"issubclass":     newBuiltinFunction("issubclass", builtinIsSubclass).ToObject(),
 		"iter":           newBuiltinFunction("iter", builtinIter).ToObject(),
 		"len":            newBuiltinFunction("len", builtinLen).ToObject(),
+		"map":            newBuiltinFunction("map", builtinMapFn).ToObject(),
 		"max":            newBuiltinFunction("max", builtinMax).ToObject(),
 		"min":            newBuiltinFunction("min", builtinMin).ToObject(),
 		"next":           newBuiltinFunction("next", builtinNext).ToObject(),
@@ -737,3 +766,57 @@ func numberToBase(prefix string, base int, o *Object) string {
 	}
 	return prefix + s
 }
+
+// initIters return list of initiated Iter instances from the list of
+// iterables.
+func initIters(f *Frame, items []*Object) ([]*Object, *BaseException) {
+	l := len(items)
+	iters := make([]*Object, l)
+	for i, arg := range items {
+		iter, raised := Iter(f, arg)
+		if raised != nil {
+			return nil, raised
+		}
+		iters[i] = iter
+	}
+	return iters, nil
+}
+
+// zipLongest return the list of aggregates elements from each of the
+// iterables. If the iterables are of uneven length, missing values are
+// filled-in with None.
+func zipLongest(f *Frame, args Args) ([][]*Object, *BaseException) {
+	argc := len(args)
+	result := make([][]*Object, 0, 2)
+	iters, raised := initIters(f, args)
+	if raised != nil {
+		return nil, raised
+	}
+
+	for {
+		noItems := true
+		elems := make([]*Object, argc)
+		for i, iter := range iters {
+			if iter == nil {
+				continue
+			}
+			elem, raised := Next(f, iter)
+			if raised != nil {
+				if raised.isInstance(StopIterationType) {
+					iters[i] = nil
+					elems[i] = None
+					continue
+				}
+				f.RestoreExc(nil, nil)
+				return nil, raised
+			}
+			noItems = false
+			elems[i] = elem
+		}
+		if noItems {
+			break
+		}
+		result = append(result, elems)
+	}
+	return result, nil
+}

runtime/builtin_types_test.go

@@ -164,6 +164,25 @@ func TestBuiltinFuncs(t *testing.T) {
 		{f: "iter", args: wrapArgs(42), wantExc: mustCreateException(TypeErrorType, "'int' object is not iterable")},
 		{f: "len", args: wrapArgs(newTestList(1, 2, 3)), want: NewInt(3).ToObject()},
 		{f: "len", args: wrapArgs(), wantExc: mustCreateException(TypeErrorType, "'len' requires 1 arguments")},
+		{f: "map", args: wrapArgs(), wantExc: mustCreateException(TypeErrorType, "map() requires at least two args")},
+		{f: "map", args: wrapArgs(StrType), wantExc: mustCreateException(TypeErrorType, "map() requires at least two args")},
+		{f: "map", args: wrapArgs(None, newTestList()), want: newTestList().ToObject()},
+		{f: "map", args: wrapArgs(None, newTestList(1, 2, 3)), want: newTestList(1, 2, 3).ToObject()},
+		{f: "map", args: wrapArgs(None, newTestDict("foo", 1, "bar", 3)), want: newTestList("foo", "bar").ToObject()},
+		{f: "map", args: wrapArgs(None, None), wantExc: mustCreateException(TypeErrorType, "'NoneType' object is not iterable")},
+		{f: "map", args: wrapArgs(StrType, None), wantExc: mustCreateException(TypeErrorType, "'NoneType' object is not iterable")},
+		{f: "map", args: wrapArgs(StrType, newTestList(), None), wantExc: mustCreateException(TypeErrorType, "'NoneType' object is not iterable")},
+		{f: "map", args: wrapArgs(newTestList(), newTestList(1, 2, 3)), wantExc: mustCreateException(TypeErrorType, "'list' object is not callable")},
+		{f: "map", args: wrapArgs(StrType, newTestList()), want: newTestList().ToObject()},
+		{f: "map", args: wrapArgs(StrType, newTestList(1, 2, 3)), want: newTestList("1", "2", "3").ToObject()},
+		{f: "map", args: wrapArgs(StrType, newTestList(-1, -2, -3)), want: newTestList("-1", "-2", "-3").ToObject()},
+		{f: "map", args: wrapArgs(IntType, newTestList("1", "2", "3")), want: newTestList(1, 2, 3).ToObject()},
+		{f: "map", args: wrapArgs(IntType, newTestList("-1", "-2", "-3")), want: newTestList(-1, -2, -3).ToObject()},
+		{f: "map", args: wrapArgs(IntType, "123"), want: newTestList(1, 2, 3).ToObject()},
+		{f: "map", args: wrapArgs(IntType, newTestDict("1", "11", "2", "22")), want: newTestList(1, 2).ToObject()},
+		{f: "map", args: wrapArgs(IntType, 1), wantExc: mustCreateException(TypeErrorType, "'int' object is not iterable")},
+		{f: "map", args: wrapArgs(1, newTestList(1, 2, 3)), wantExc: mustCreateException(TypeErrorType, "'int' object is not callable")},
+		{f: "map", args: wrapArgs(StrType, newTestList(), 1), wantExc: mustCreateException(TypeErrorType, "'int' object is not iterable")},
 		{f: "max", args: wrapArgs(2, 3, 1), want: NewInt(3).ToObject()},
 		{f: "max", args: wrapArgs("bar", "foo"), want: NewStr("foo").ToObject()},
 		{f: "max", args: wrapArgs(newTestList(2, 3, 1)), want: NewInt(3).ToObject()},

testing/builtin_test.py

@@ -281,3 +281,23 @@ try:
   zip([1, 2, 3], [1, 2], [4], None)
 except TypeError as e:
   assert str(e) == "'NoneType' object is not iterable"
+
+# Test map
+
+assert map(str, []) == []
+assert map(str, [1, 2, 3]) == ["1", "2", "3"]
+assert map(str, (1, 2, 3)) == ["1", "2", "3"]
+# assert map(str, (1.0, 2.0, 3.0)) == ["1", "2", "3"]
+assert map(str, range(3)) == ["0", "1", "2"]
+assert map(str, xrange(3)) == ["0", "1", "2"]
+assert map(int, ["1", "2", "3"]) == [1, 2, 3]
+assert map(int, "123") == [1, 2, 3]
+assert map(int, {"1": "a", "2": "b"}) == [1, 2]
+assert map(int, {1: "a", 2: "b"}) == [1, 2]
+assert map(lambda a, b: (str(a), float(b or 0) + 0.1),
+           [1, 2, 3], [1, 2]) == [('1', 1.1), ('2', 2.1), ('3', 0.1)]
+assert map(None, [1, 2, 3]) == [1, 2, 3]
+a = [1, 2, 3]
+assert map(None, a) == a
+assert map(None, a) is not a
+assert map(None, (1, 2, 3)) == [1, 2, 3]