From dc5c18517725f2c885531c61ff7d44f1f6427146 Mon Sep 17 00:00:00 2001 From: Ang Long Date: Sun, 30 Jul 2017 23:59:11 +0800 Subject: [PATCH] Add builtin filter function --- runtime/builtin_types.go | 81 +++++++++++++++++++++++++++++++++++ runtime/builtin_types_test.go | 38 ++++++++++++++++ 2 files changed, 119 insertions(+) diff --git a/runtime/builtin_types.go b/runtime/builtin_types.go index 888951e7..5075a700 100644 --- a/runtime/builtin_types.go +++ b/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(), diff --git a/runtime/builtin_types_test.go b/runtime/builtin_types_test.go index 732db19d..00d9da22 100644 --- a/runtime/builtin_types_test.go +++ b/runtime/builtin_types_test.go @@ -98,6 +98,34 @@ 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(), + })) + customIterStrType := newTestClass("CustomIterStrType", []*Type{StrType}, newStringDict(map[string]*Object{ + "__iter__": newBuiltinFunction("__iter__", func(f *Frame, args Args, kwargs KWArgs) (*Object, *BaseException) { + return newObject(createNextType(3)), nil + }).ToObject(), + })) + customIterTupleType := newTestClass("CustomIterTupleType", []*Type{TupleType}, 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 +214,16 @@ 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: "filter", args: wrapArgs(BoolType, &Str{Object: Object{typ: customIterStrType}, value: "foo"}), want: NewStr("foo").ToObject()}, + {f: "filter", args: wrapArgs(BoolType, &Tuple{Object: Object{typ: customIterTupleType}, elems: []*Object{NewInt(5).ToObject()}}), want: NewTuple1(NewInt(5).ToObject()).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()},