Add method to build string query and method to update tree by some rules method

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()

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

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