From a3a22f545e149b714b6f2b17cd04a5d6e0bd5294 Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Fri, 23 Sep 2022 19:07:13 +0200 Subject: Adding upstream version 6.2.0. Signed-off-by: Daniel Baumann --- tests/test_expressions.py | 32 +++++++++++++++++++++++++++++--- 1 file changed, 29 insertions(+), 3 deletions(-) (limited to 'tests/test_expressions.py') diff --git a/tests/test_expressions.py b/tests/test_expressions.py index 716e457..59d584c 100644 --- a/tests/test_expressions.py +++ b/tests/test_expressions.py @@ -224,9 +224,6 @@ class TestExpressions(unittest.TestCase): self.assertEqual(actual_expression_2.sql(dialect="presto"), "IF(c - 2 > 0, c - 2, b)") self.assertIs(actual_expression_2, expression) - with self.assertRaises(ValueError): - parse_one("a").transform(lambda n: None) - def test_transform_no_infinite_recursion(self): expression = parse_one("a") @@ -247,6 +244,35 @@ class TestExpressions(unittest.TestCase): self.assertEqual(expression.transform(fun).sql(), "SELECT a, b FROM x") + def test_transform_node_removal(self): + expression = parse_one("SELECT a, b FROM x") + + def remove_column_b(node): + if isinstance(node, exp.Column) and node.name == "b": + return None + return node + + self.assertEqual(expression.transform(remove_column_b).sql(), "SELECT a FROM x") + self.assertEqual(expression.transform(lambda _: None), None) + + expression = parse_one("CAST(x AS FLOAT)") + + def remove_non_list_arg(node): + if isinstance(node, exp.DataType): + return None + return node + + self.assertEqual(expression.transform(remove_non_list_arg).sql(), "CAST(x AS )") + + expression = parse_one("SELECT a, b FROM x") + + def remove_all_columns(node): + if isinstance(node, exp.Column): + return None + return node + + self.assertEqual(expression.transform(remove_all_columns).sql(), "SELECT FROM x") + def test_replace(self): expression = parse_one("SELECT a, b FROM x") expression.find(exp.Column).replace(parse_one("c")) -- cgit v1.2.3