def test_mapping_complex_expression() -> None: """ Maps over an Expression container: f0(t1.c1, fB(f())) -> f0(t1.c1, fB(f(f() as a))) """ f5 = FunctionCall("a", "f", ()) f4 = FunctionCall(None, "f", (f5, )) f3 = FunctionCall(None, "f", ()) def replace_expr(e: Expression) -> Expression: if isinstance(e, FunctionCall) and e.function_name == "f": return f4 return e c1 = Column(None, "c1", "t1") f2 = FunctionCall(None, "fB", (f3, )) f1 = FunctionCall(None, "f0", (c1, f2)) # Only the external function is going to be replaced since, when map returns a new # column, we expect the func to have takern care of its own children. f1 = f1.transform(replace_expr) iterate = list(f1) expected = [ c1, f5, f4, FunctionCall(None, "fB", (f4, )), FunctionCall(None, "f0", (c1, FunctionCall(None, "fB", (f4, )))), ] assert iterate == expected
def test_add_alias() -> None: """ Adds an alias to a column referenced in a function f(t1.c1) -> f(t1.c1 as a) """ column1 = Column(None, "c1", "t1") column2 = Column("a", "c1", "t1") def replace_expr(e: Expression) -> Expression: if isinstance(e, Column) and e.column_name == "c1": return column2 return e f = FunctionCall(None, "f", (column1, )) f2 = f.transform(replace_expr) expected = [column2, FunctionCall(None, "f", (column2, ))] assert list(f2) == expected