Skip to content
This repository was archived by the owner on Mar 23, 2023. It is now read-only.

Add builtin filter function #374

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add builtin filter function #374

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 81 additions & 0 deletions runtime/builtin_types.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
package grumpy

import (
"bytes"
"fmt"
"math"
"math/big"
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(),
Expand Down
38 changes: 38 additions & 0 deletions runtime/builtin_types_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()},
Expand Down