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

Commit dc5c185

Browse files
aiskaisk
committed
Add builtin filter function
1 parent 6bdfc4b commit dc5c185

File tree

2 files changed

+119
-0
lines changed

2 files changed

+119
-0
lines changed

‎runtime/builtin_types.go‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
package grumpy
1616

1717
import (
18+
"bytes"
1819
"fmt"
1920
"math"
2021
"math/big"
@@ -248,6 +249,85 @@ func builtinMapFn(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
248249
return NewList(result...).ToObject(), nil
249250
}
250251

252+
func builtinFilter(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
253+
if raised := checkMethodArgs(f, "filter", args, ObjectType, ObjectType); raised != nil {
254+
return nil, raised
255+
}
256+
fn := args[0]
257+
l := args[1]
258+
filterFunc := IsTrue
259+
if fn != None {
260+
filterFunc = func(f *Frame, o *Object) (bool, *BaseException) {
261+
result, raised := fn.Call(f, Args{o}, nil)
262+
if raised != nil {
263+
return false, raised
264+
}
265+
return IsTrue(f, result)
266+
}
267+
}
268+
switch {
269+
// CPython will return the same type if the second type is tuple or string, else return a list.
270+
case l.isInstance(TupleType):
271+
result := make([]*Object, 0)
272+
for _, item := range toTupleUnsafe(l).elems {
273+
ret, raised := filterFunc(f, item)
274+
if raised != nil {
275+
return nil, raised
276+
}
277+
if ret {
278+
result = append(result, item)
279+
}
280+
}
281+
return NewTuple(result...).ToObject(), nil
282+
case l.isInstance(StrType):
283+
if fn == None {
284+
return l, nil
285+
}
286+
var result bytes.Buffer
287+
for _, item := range []byte(toStrUnsafe(l).Value()) {
288+
ret, raised := filterFunc(f, NewStr(string(item)).ToObject())
289+
if raised != nil {
290+
return nil, raised
291+
}
292+
if ret {
293+
result.WriteByte(item)
294+
}
295+
}
296+
return NewStr(result.String()).ToObject(), nil
297+
case l.isInstance(UnicodeType):
298+
if fn == None {
299+
return l, nil
300+
}
301+
var result []rune
302+
for _, item := range toUnicodeUnsafe(l).Value() {
303+
ret, raised := filterFunc(f, NewUnicodeFromRunes([]rune{item}).ToObject())
304+
if raised != nil {
305+
return nil, raised
306+
}
307+
if ret {
308+
result = append(result, item)
309+
}
310+
}
311+
return NewUnicodeFromRunes(result).ToObject(), nil
312+
default:
313+
result := make([]*Object, 0)
314+
raised := seqForEach(f, l, func(item *Object) (raised *BaseException) {
315+
ret, raised := filterFunc(f, item)
316+
if raised != nil {
317+
return raised
318+
}
319+
if ret {
320+
result = append(result, item)
321+
}
322+
return nil
323+
})
324+
if raised != nil {
325+
return nil, raised
326+
}
327+
return NewList(result...).ToObject(), nil
328+
}
329+
}
330+
251331
func builtinAll(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
252332
if raised := checkFunctionArgs(f, "all", args, ObjectType); raised != nil {
253333
return nil, raised
@@ -768,6 +848,7 @@ func init() {
768848
"divmod": newBuiltinFunction("divmod", builtinDivMod).ToObject(),
769849
"Ellipsis": Ellipsis,
770850
"False": False.ToObject(),
851+
"filter": newBuiltinFunction("filter", builtinFilter).ToObject(),
771852
"getattr": newBuiltinFunction("getattr", builtinGetAttr).ToObject(),
772853
"globals": newBuiltinFunction("globals", builtinGlobals).ToObject(),
773854
"hasattr": newBuiltinFunction("hasattr", builtinHasAttr).ToObject(),

‎runtime/builtin_types_test.go‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,34 @@ func TestBuiltinFuncs(t *testing.T) {
9898
return newObject(badNextType), nil
9999
}).ToObject(),
100100
}))
101+
createNextType := func(n int) *Type {
102+
i := -1
103+
nextType := newTestClass("FooNextType", []*Type{ObjectType}, newStringDict(map[string]*Object{
104+
"next": newBuiltinFunction("next", func(f *Frame, args Args, kwargs KWArgs) (*Object, *BaseException) {
105+
if i >= n {
106+
return nil, f.RaiseType(StopIterationType, "foo")
107+
}
108+
i++
109+
return NewInt(i).ToObject(), nil
110+
}).ToObject(),
111+
}))
112+
return nextType
113+
}
114+
fooIterType := newTestClass("FooIterType", []*Type{ObjectType}, newStringDict(map[string]*Object{
115+
"__iter__": newBuiltinFunction("__iter__", func(f *Frame, args Args, kwargs KWArgs) (*Object, *BaseException) {
116+
return newObject(createNextType(3)), nil
117+
}).ToObject(),
118+
}))
119+
customIterStrType := newTestClass("CustomIterStrType", []*Type{StrType}, newStringDict(map[string]*Object{
120+
"__iter__": newBuiltinFunction("__iter__", func(f *Frame, args Args, kwargs KWArgs) (*Object, *BaseException) {
121+
return newObject(createNextType(3)), nil
122+
}).ToObject(),
123+
}))
124+
customIterTupleType := newTestClass("CustomIterTupleType", []*Type{TupleType}, newStringDict(map[string]*Object{
125+
"__iter__": newBuiltinFunction("__iter__", func(f *Frame, args Args, kwargs KWArgs) (*Object, *BaseException) {
126+
return newObject(createNextType(3)), nil
127+
}).ToObject(),
128+
}))
101129
addType := newTestClass("Add", []*Type{ObjectType}, newStringDict(map[string]*Object{
102130
"__add__": newBuiltinFunction("__add__", func(f *Frame, _ Args, _ KWArgs) (*Object, *BaseException) {
103131
return NewInt(1).ToObject(), nil
@@ -186,6 +214,16 @@ func TestBuiltinFuncs(t *testing.T) {
186214
{f: "divmod", args: wrapArgs(-3.25, -1.0), want: NewTuple2(NewFloat(3.0).ToObject(), NewFloat(-0.25).ToObject()).ToObject()},
187215
{f: "divmod", args: wrapArgs(NewStr("a"), NewStr("b")), wantExc: mustCreateException(TypeErrorType, "unsupported operand type(s) for divmod(): 'str' and 'str'")},
188216
{f: "divmod", args: wrapArgs(), wantExc: mustCreateException(TypeErrorType, "'divmod' requires 2 arguments")},
217+
{f: "filter", args: wrapArgs(None, NewTuple2(NewInt(0).ToObject(), NewInt(1).ToObject())), want: NewTuple1(NewInt(1).ToObject()).ToObject()},
218+
{f: "filter", args: wrapArgs(BoolType, NewTuple2(NewInt(0).ToObject(), NewInt(1).ToObject())), want: NewTuple1(NewInt(1).ToObject()).ToObject()},
219+
{f: "filter", args: wrapArgs(None, "012"), want: NewStr("012").ToObject()},
220+
{f: "filter", args: wrapArgs(IntType, "012"), want: NewStr("12").ToObject()},
221+
{f: "filter", args: wrapArgs(None, NewUnicode("012")), want: NewUnicode("012").ToObject()},
222+
{f: "filter", args: wrapArgs(None, newTestList(1, 0, 3)), want: newTestList(1, 3).ToObject()},
223+
{f: "filter", args: wrapArgs(IntType, newTestList("1", "0", "3")), want: newTestList("1", "3").ToObject()},
224+
{f: "filter", args: wrapArgs(BoolType, newObject(fooIterType)), want: newTestList(1, 2, 3).ToObject()},
225+
{f: "filter", args: wrapArgs(BoolType, &Str{Object: Object{typ: customIterStrType}, value: "foo"}), want: NewStr("foo").ToObject()},
226+
{f: "filter", args: wrapArgs(BoolType, &Tuple{Object: Object{typ: customIterTupleType}, elems: []*Object{NewInt(5).ToObject()}}), want: NewTuple1(NewInt(5).ToObject()).ToObject()},
189227
{f: "getattr", args: wrapArgs(None, NewStr("foo").ToObject(), NewStr("bar").ToObject()), want: NewStr("bar").ToObject()},
190228
{f: "getattr", args: wrapArgs(None, NewStr("foo").ToObject()), wantExc: mustCreateException(AttributeErrorType, "'NoneType' object has no attribute 'foo'")},
191229
{f: "hasattr", args: wrapArgs(newObject(ObjectType), NewStr("foo").ToObject()), want: False.ToObject()},

0 commit comments

Comments
 (0)