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
