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,85 @@ func builtinMapFn(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
 	return NewList(result...).ToObject(), nil
 }
 
+func builtinFilter(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
+	if raised := checkMethodArgs(f, "filter", args, ObjectType, ObjectType); raised != nil {
+		return nil, raised
+	}
+	fn := args[0]
+	l := args[1]
+	filterFunc := IsTrue
+	if fn != None {
+		filterFunc = func(f *Frame, o *Object) (bool, *BaseException) {
+			result, raised := fn.Call(f, Args{o}, nil)
+			if raised != nil {
+				return false, raised
+			}
+			return IsTrue(f, result)
+		}
+	}
+	switch {
+	// CPython will return the same type if the second type is tuple or string, else return a list.
+	case l.isInstance(TupleType):
+		result := make([]*Object, 0)
+		for _, item := range toTupleUnsafe(l).elems {
+			ret, raised := filterFunc(f, item)
+			if raised != nil {
+				return nil, raised
+			}
+			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 []byte(toStrUnsafe(l).Value()) {
+			ret, raised := filterFunc(f, NewStr(string(item)).ToObject())
+			if raised != nil {
+				return nil, raised
+			}
+			if ret {
+				result.WriteByte(item)
+			}
+		}
+		return NewStr(result.String()).ToObject(), nil
+	case l.isInstance(UnicodeType):
+		if fn == None {
+			return l, nil
+		}
+		var result []rune
+		for _, item := range toUnicodeUnsafe(l).Value() {
+			ret, raised := filterFunc(f, NewUnicodeFromRunes([]rune{item}).ToObject())
+			if raised != nil {
+				return nil, raised
+			}
+			if ret {
+				result = append(result, item)
+			}
+		}
+		return NewUnicodeFromRunes(result).ToObject(), nil
+	default:
+		result := make([]*Object, 0)
+		raised := seqForEach(f, l, func(item *Object) (raised *BaseException) {
+			ret, raised := filterFunc(f, item)
+			if raised != nil {
+				return raised
+			}
+			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 +848,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

@@ -98,6 +98,24 @@ func TestBuiltinFuncs(t *testing.T) {
 			return newObject(badNextType), nil
 		}).ToObject(),
 	}))
+	createNextType := func(n int) *Type {
+		i := -1
+		nextType := newTestClass("FooNextType", []*Type{ObjectType}, newStringDict(map[string]*Object{
+			"next": newBuiltinFunction("next", func(f *Frame, args Args, kwargs KWArgs) (*Object, *BaseException) {
+				if i >= n {
+					return nil, f.RaiseType(StopIterationType, "foo")
+				}
+				i++
+				return NewInt(i).ToObject(), nil
+			}).ToObject(),
+		}))
+		return nextType
+	}
+	fooIterType := newTestClass("FooIterType", []*Type{ObjectType}, newStringDict(map[string]*Object{
+		"__iter__": newBuiltinFunction("__iter__", func(f *Frame, args Args, kwargs KWArgs) (*Object, *BaseException) {
+			return newObject(createNextType(3)), nil
+		}).ToObject(),
+	}))
 	addType := newTestClass("Add", []*Type{ObjectType}, newStringDict(map[string]*Object{
 		"__add__": newBuiltinFunction("__add__", func(f *Frame, _ Args, _ KWArgs) (*Object, *BaseException) {
 			return NewInt(1).ToObject(), nil
@@ -186,6 +204,14 @@ 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(None, "012"), want: NewStr("012").ToObject()},
+		{f: "filter", args: wrapArgs(IntType, "012"), want: NewStr("12").ToObject()},
+		{f: "filter", args: wrapArgs(None, NewUnicode("012")), want: NewUnicode("012").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: "filter", args: wrapArgs(BoolType, newObject(fooIterType)), want: newTestList(1, 2, 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()},