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