diff --git a/lucyparser/tree.py b/lucyparser/tree.py index 1b9dcc2..57d71fa 100644 --- a/lucyparser/tree.py +++ b/lucyparser/tree.py @@ -38,6 +38,9 @@ class RawOperator: } +OPERATOR_TO_RAW_OPERATOR = {v: k for k, v in RAW_OPERATOR_TO_OPERATOR.items()} + + class LogicalOperator(enum.Enum): NOT = enum.auto() AND = enum.auto() diff --git a/lucyparser/utils.py b/lucyparser/utils.py new file mode 100644 index 0000000..907d88c --- /dev/null +++ b/lucyparser/utils.py @@ -0,0 +1,51 @@ +from typing import Callable + +from lucyparser.tree import BaseNode, ExpressionNode, OPERATOR_TO_RAW_OPERATOR, Operator, OrNode, AndNode, NotNode + + +def build_query(tree: BaseNode, level=0) -> str: + """ + :param tree: Parsed tree + :param level: Current tree level + :return: query as string + """ + if isinstance(tree, ExpressionNode): + raw_op = OPERATOR_TO_RAW_OPERATOR[tree.operator] + operator = raw_op if tree.operator == Operator.EQ else f" {raw_op}" + + value = tree.value + if isinstance(value, str) and '"' in value: + value = value.replace('"', '\\"') + + return f"""{tree.name}{operator} \"{value}\"""" + + format_str = "{}" if not level else "({})" + if isinstance(tree, OrNode): + return format_str.format(" OR ".join(build_query(child, level + 1) for child in tree.children)) + if isinstance(tree, AndNode): + return format_str.format(" AND ".join(build_query(child, level + 1) for child in tree.children)) + + if isinstance(tree, NotNode): + return "NOT {}".format(build_query(tree.children[0], level + 1)) + + +def update_tree(tree: BaseNode, handler: Callable[[BaseNode], int]) -> int: + """ + :param tree: Parsed tree + :param handler: Function with ExpressionNode as argument. + That function is able to change current node or do nothing. + It must return replacement count + :return: replacement count + """ + if isinstance(tree, ExpressionNode): + return handler(tree) + + elif isinstance(tree, (AndNode, OrNode, NotNode)): + replacement_count = handler(tree) + + for child in tree.children: + replacement_count += update_tree(tree=child, handler=handler) + + return replacement_count + + return 0 diff --git a/tests/test_parsing.py b/tests/test_parsing.py index 6384d6f..b17a7f9 100644 --- a/tests/test_parsing.py +++ b/tests/test_parsing.py @@ -2,7 +2,8 @@ from lucyparser import parse from lucyparser.parsing import Cursor -from lucyparser.tree import ExpressionNode, Operator, NotNode, AndNode, OrNode +from lucyparser.tree import ExpressionNode, Operator, NotNode, AndNode, OrNode, BaseNode +from lucyparser.utils import build_query, update_tree @pytest.mark.parametrize( @@ -99,3 +100,49 @@ def test_simple_case(raw, parsed): ) def test_starts_with_a_word(string, word, result): assert Cursor(string).starts_with_a_word(word) is result + + +@pytest.mark.parametrize( + "query,expected", + [ + ('x: y OR y: z AND (NOT n: n OR ((((l: z)))))', 'x: "y" OR (y: "z" AND (NOT n: "n" OR l: "z"))'), + ('(((NOT (x: y OR b:z))))', 'NOT (x: "y" OR b: "z")'), + ('(((NOT (x: y) OR b:z)))', 'NOT x: "y" OR b: "z"'), + ('(((NOT (x: y) OR b:"\\"ululul")))', 'NOT x: "y" OR b: "\\"ululul"'), + ('NOT x: y', 'NOT x: "y"'), + ], +) +def test_get_str_from_tree(query, expected): + assert build_query(parse(query)) == expected + + +@pytest.mark.parametrize( + "query,replaced_value_x_query,replaced_key_x_query", + [ + ('x: "y"', 'x: "y"', 'not_x: "y"'), + ('y: "x"', 'y: "not_x"', 'y: "x"'), + ('x: "x"', 'x: "not_x"', 'not_x: "x"'), + ], +) +def test_update_tree(query, replaced_value_x_query, replaced_key_x_query): + def replace_value_x(tree: BaseNode) -> int: + if isinstance(tree, ExpressionNode) and tree.value == "x": + tree.value = "not_x" + return 1 + return 0 + + # replace value x + tree = parse(query) + update_tree(tree=tree, handler=replace_value_x) + assert build_query(tree) == replaced_value_x_query + + def replace_key_x(tree: BaseNode) -> int: + if isinstance(tree, ExpressionNode) and tree.name == "x": + tree.name = "not_x" + return 1 + return 0 + + # replace key x + tree = parse(query) + update_tree(tree=tree, handler=replace_key_x) + assert build_query(tree) == replaced_key_x_query