def __call__(self, *args, **kwargs):
call = Call('__call__',
BinaryOp('__getattr__', MetaArg("_"), self.name),
*args,
**kwargs
)
return self.to_copy(name = call)
def test_call_tree_local_dispatch_fail(f_dispatch_strict):
ctl = CallTreeLocal({'f_a': lambda self: self}, dispatch_cls=object)
call = Call("__call__", FuncArg(f_dispatch_strict), MetaArg('_'))
# should be the default failure dispatch for object
new_call = ctl.enter(call)
with pytest.raises(TypeError):
new_call('na')
def top_n(__data, n, wt=None):
"""Filter to keep the top or bottom entries in each group.
Args:
___data: a DataFrame
n: the number of rows to keep in each group
wt: a column or expression that determines ordering (defaults to the last column in data)
Examples:
>>> from siuba import _, top_n
>>> df = pd.DataFrame({'x': [3, 1, 2, 4], 'y': [1, 1, 0, 0]})
>>> top_n(df, 2, _.x)
x y
0 3 1
3 4 0
>>> top_n(df, -2, _.x)
x y
1 1 1
2 2 0
>>> top_n(df, 2, _.x*_.y)
x y
0 3 1
1 1 1
"""
# NOTE: using min_rank, since it can return a lazy expr for min_rank(ing)
# but I would rather not have it imported in verbs. will be more
# reasonable if each verb were its own file? need abstract verb / vector module.
# vector imports experimental right now, so we need to invert deps
# TODO:
# * what if wt is a string? should remove all str -> expr in verbs like group_by etc..
# * verbs like filter allow lambdas, but this func breaks with that
from .vector import min_rank
if wt is None:
sym_wt = getattr(Symbolic(MetaArg("_")), __data.columns[-1])
elif isinstance(wt, Call):
sym_wt = Symbolic(wt)
else:
raise TypeError("wt must be a symbolic expression, eg. _.some_col")
if n > 0:
return filter(__data, min_rank(-sym_wt) <= n)
else:
return filter(__data, min_rank(sym_wt) <= abs(n))
def pipe_with_meta(f, *args, **kwargs):
return create_pipe_call(f, MetaArg("_"), *args, **kwargs)
def wrapper(*args, **kwargs):
return create_pipe_call(f, MetaArg("_"), *args, **kwargs)
def test_call_tree_local_dispatch_cls_subclass(f_dispatch):
ctl = CallTreeLocal({'f_a': lambda self: self}, dispatch_cls=SomeClass)
call = Call("__call__", FuncArg(f_dispatch), MetaArg('_'))
new_call = ctl.enter(call)
assert new_call('na') == 'some class'
def test_call_tree_local_dispatch_cls_object(f_dispatch):
ctl = CallTreeLocal({'f_a': lambda self: self}, dispatch_cls=object)
call = Call("__call__", FuncArg(f_dispatch), MetaArg('_'))
new_call = ctl.enter(call)
assert new_call('na') == 'default'
res = slice_call(data)
assert res == target
# Node copying ================================================================
from siuba.siu import Call, BinaryOp, SliceOp, MetaArg, FuncArg, DictCall
# Call
@pytest.mark.parametrize('node', [
Call("__call__", lambda x, y=2: x + y, 1, y=2),
BinaryOp("__add__", 1, 2),
SliceOp("__siu_slice__", slice(0, 1)),
SliceOp("__siu_slice__", (slice(0, 1), slice(2, 3))),
MetaArg("_"),
FuncArg("__custom_func__", lambda x: x),
FuncArg(lambda x: x),
DictCall("__call__", dict, {
'a': 1,
'b': 2
})
])
def test_node_copy(node):
copy = node.copy()
assert isinstance(copy, node.__class__)
assert copy is not node
assert copy.func == node.func
assert copy.args == node.args
assert copy.kwargs == node.kwargs
