Add builtin filter function

Author: aisk

runtime/builtin_types.go

@@ -15,6 +15,7 @@
 package grumpy
 
 import (
+	"bytes"
 	"fmt"
 	"math"
 	"math/big"
@@ -248,6 +249,82 @@ func builtinMapFn(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
 	return NewList(result...).ToObject(), nil
 }
 
+func builtinFilter(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
+	argc := len(args)
+	if argc < 2 {
+		return nil, f.RaiseType(TypeErrorType, fmt.Sprintf("filter expected 2 arguments, got %d", argc))
+	}
+	fn := args[0]
+	l := args[1]
+	switch {
+	case l.isInstance(TupleType):
+		result := make([]*Object, 0)
+		for _, item := range toTupleUnsafe(l).elems {
+			if fn == None {
+				if ret, raised := IsTrue(f, item); raised != nil {
+					return nil, raised
+				} else if ret {
+					result = append(result, item)
+				}
+			} else {
+				ret, raised := args[0].Call(f, []*Object{item}, nil)
+				if raised != nil {
+					return nil, raised
+				}
+				if ret, raised := IsTrue(f, ret); raised != nil {
+					return nil, raised
+				} else if ret {
+					result = append(result, item)
+				}
+			}
+		}
+		return NewTuple(result...).ToObject(), nil
+	case l.isInstance(StrType):
+		if fn == None {
+			return l, nil
+		}
+		var result bytes.Buffer
+		for _, item := range toStrUnsafe(l).Value() {
+			ret, raised := fn.Call(f, []*Object{NewStr(string(item)).ToObject()}, nil)
+			if raised != nil {
+				return nil, raised
+			}
+			if ret, raised := IsTrue(f, ret); raised != nil {
+				return nil, raised
+			} else if ret {
+				result.WriteRune(item)
+			}
+		}
+		return NewStr(result.String()).ToObject(), nil
+	default:
+		result := make([]*Object, 0)
+		raised := seqForEach(f, l, func(item *Object) (raised *BaseException) {
+			if fn == None {
+				if ret, raised := IsTrue(f, item); raised != nil {
+					return raised
+				} else if ret {
+					result = append(result, item)
+				}
+			} else {
+				ret, raised := fn.Call(f, []*Object{item}, nil)
+				if raised != nil {
+					return raised
+				}
+				if ret, raised := IsTrue(f, ret); raised != nil {
+					return raised
+				} else if ret {
+					result = append(result, item)
+				}
+			}
+			return nil
+		})
+		if raised != nil {
+			return nil, raised
+		}
+		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
@@ -768,6 +845,7 @@ func init() {
 		"divmod":         newBuiltinFunction("divmod", builtinDivMod).ToObject(),
 		"Ellipsis":       Ellipsis,
 		"False":          False.ToObject(),
+		"filter":         newBuiltinFunction("filter", builtinFilter).ToObject(),
 		"getattr":        newBuiltinFunction("getattr", builtinGetAttr).ToObject(),
 		"globals":        newBuiltinFunction("globals", builtinGlobals).ToObject(),
 		"hasattr":        newBuiltinFunction("hasattr", builtinHasAttr).ToObject(),

runtime/builtin_types_test.go

@@ -186,6 +186,11 @@ func TestBuiltinFuncs(t *testing.T) {
 		{f: "divmod", args: wrapArgs(-3.25, -1.0), want: NewTuple2(NewFloat(3.0).ToObject(), NewFloat(-0.25).ToObject()).ToObject()},
 		{f: "divmod", args: wrapArgs(NewStr("a"), NewStr("b")), wantExc: mustCreateException(TypeErrorType, "unsupported operand type(s) for divmod(): 'str' and 'str'")},
 		{f: "divmod", args: wrapArgs(), wantExc: mustCreateException(TypeErrorType, "'divmod' requires 2 arguments")},
+		{f: "filter", args: wrapArgs(None, NewTuple2(NewInt(0).ToObject(), NewInt(1).ToObject())), want: NewTuple1(NewInt(1).ToObject()).ToObject()},
+		{f: "filter", args: wrapArgs(BoolType, NewTuple2(NewInt(0).ToObject(), NewInt(1).ToObject())), want: NewTuple1(NewInt(1).ToObject()).ToObject()},
+		{f: "filter", args: wrapArgs(IntType, "012"), want: NewStr("12").ToObject()},
+		{f: "filter", args: wrapArgs(None, newTestList(1, 0, 3)), want: newTestList(1, 3).ToObject()},
+		{f: "filter", args: wrapArgs(IntType, newTestList("1", "0", "3")), want: newTestList("1", "3").ToObject()},
 		{f: "getattr", args: wrapArgs(None, NewStr("foo").ToObject(), NewStr("bar").ToObject()), want: NewStr("bar").ToObject()},
 		{f: "getattr", args: wrapArgs(None, NewStr("foo").ToObject()), wantExc: mustCreateException(AttributeErrorType, "'NoneType' object has no attribute 'foo'")},
 		{f: "hasattr", args: wrapArgs(newObject(ObjectType), NewStr("foo").ToObject()), want: False.ToObject()},