def test_lambda_in_filter(): df = DataFrame() df1 = df[df.apply(lambda e: e == 1)] assert isinstance(df1.child_expr, ast_DataFrame) assert df1.filter is not None assert isinstance(df1.filter, Column) assert isinstance(df1.filter.child_expr, ast.Call)
def test_lambda_argument(): df = DataFrame() df1 = df.apply(lambda e: e) assert df1.child_expr is not None assert isinstance(df1.child_expr, ast.Call) assert len(df1.child_expr.args) == 1 arg1 = df1.child_expr.args[0] assert isinstance(arg1, ast_Callable)
def test_callable_wrong_number_args(): d = DataFrame() d1 = d.apply(lambda b: b) expr, ctx = render(d1) assert isinstance(expr, ast.Call) arg1 = expr.args[0] # type: ast.AST assert isinstance(arg1, ast_Callable) with pytest.raises(Exception): render_callable(arg1, ctx, d, d)
def test_callable_returns_const(): d = DataFrame() d1 = d.apply(lambda b: 20) expr, ctx = render(d1) assert isinstance(expr, ast.Call) arg1 = expr.args[0] # type: ast.AST assert isinstance(arg1, ast_Callable) expr1, new_ctx = render_callable(arg1, ctx, d) assert isinstance(expr1, ast.Num) assert expr1.n == 20
def test_callable_simple_call(): d = DataFrame() d1 = d.apply(lambda b: b) expr, ctx = render(d1) assert isinstance(expr, ast.Call) arg1 = expr.args[0] # type: ast.AST assert isinstance(arg1, ast_Callable) expr1, new_ctx = render_callable(arg1, ctx, d) assert isinstance(expr1, ast_DataFrame) assert expr1.dataframe is d
def test_callable_function(): def test_func(b): return b d = DataFrame() d1 = d.apply(test_func) expr, ctx = render(d1) assert isinstance(expr, ast.Call) arg1 = expr.args[0] # type: ast.AST assert isinstance(arg1, ast_Callable) expr1, new_ctx = render_callable(arg1, ctx, d) assert isinstance(expr1, ast_DataFrame) assert expr1.dataframe is d
def test_callable_context_no_update(): d = DataFrame() d1 = d.apply(lambda b: b.jets.pt) expr, ctx = render(d1) assert isinstance(expr, ast.Call) arg1 = expr.args[0] # type: ast.AST assert isinstance(arg1, ast_Callable) seen_ds = len(ctx._seen_datasources) resolved = len(ctx._resolved) expr1, new_ctx = render_callable(arg1, ctx, d) assert len(new_ctx._seen_datasources) != len(ctx._seen_datasources) \ or len(new_ctx._resolved) != len(ctx._resolved) assert seen_ds == len(ctx._seen_datasources) \ or resolved == len(ctx._resolved)