def test_dict_call_child_copy():
bin_op = BinaryOp('__add__', 1, 2)
call = DictCall("__call__", dict, {'a': bin_op, 'b': 2})
copy = call.copy()
assert isinstance(call.args[1]['a'], BinaryOp)
assert bin_op is not copy.args[1]['a']
def shift_add_ops(call):
if not (isinstance(call, BinaryOp) and call.func == "__add__"):
return call
pipe = unfurl_pipe(call)
pipe_to_binary(pipe, "__add__")
if not (isinstance(pipe[0], BinaryOp) and pipe[0].func == "__rshift__"):
return call
phead, ptail = pipe[0].args
# prepend tail to add ops
add_ops = BinaryOp("__add__", ptail, pipe_to_binary(pipe[1:], "__add__"))
# re-connect head and tail
return BinaryOp("__rshift__", phead, add_ops)
def __call__(self, *args, **kwargs):
call = Call('__call__',
BinaryOp('__getattr__', MetaArg("_"), self.name),
*args,
**kwargs
)
return self.to_copy(name = call)
def call_factory(func, *args, **kwargs):
method = flip_get(BINARY_OPS, func)
if method:
return BinaryOp(method, *args, **kwargs)
method = flip_get(UNARY_OPS, func)
if method:
return UnaryOp(method, *args, **kwargs)
return Call("__call__", NameArg(func), *args, **kwargs)
def pipe_to_binary(pipe, op):
import functools
return functools.reduce(lambda acc, v: BinaryOp(op, acc, v), pipe)
sym = eval(expr, {"_": _})
index_op, slice_call = strip_symbolic(sym).args
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
def create_column_access(obj, attr):
# would be useful to move into siu
if name_well_formatted(attr):
return BinaryOp("__getattr__", obj, attr)
return BinaryOp("__getitem__", obj, attr)
