better_parse_tree_to_str.py

from typing import Dict, Any


SAMPLE_TREE = {
    "type": "DIVISION",
    "left": {
        "type": "PAREN",
        "expression": {
            "type": "ADDITION",
            "left": {"type": "VARIABLE", "name": "$b"},
            "right": {
                "type": "FUNCTION",
                "name": "SQRT",
                "arguments": [
                    {
                        "type": "SUBTRACTION",
                        "left": {
                            "type": "FUNCTION",
                            "name": "SQR",
                            "arguments": [{"type": "VARIABLE", "name": "$b"}],
                        },
                        "right": {
                            "type": "MULTIPLICATION",
                            "left": {"type": "NUMBER", "value": 4},
                            "right": {"type": "VARIABLE", "name": "$a"},
                        },
                    }
                ],
            },
        },
    },
    "right": {
        "type": "PAREN",
        "expression": {
            "type": "MULTIPLICATION",
            "left": {"type": "NUMBER", "value": 2},
            "right": {"type": "VARIABLE", "name": "$a"},
        },
    },
}


Node = Dict[str, Any]


def parse_parenthesis(node: Node) -> str:
    formula = parse_ast_node_to_formula(node["expression"])
    return f"({formula})"


def parse_function(node: Node) -> str:
    function_name = node["name"]
    argument_formulas = []
    for argument in node["arguments"]:
        argument_formulas.append(parse_ast_node_to_formula(argument))
    args = ", ".join(argument_formulas)
    return f"{function_name}({args})"


# Map a "type" to a function to handle Node -> str transformation!
SERIALIZATION_FUNCTIONS = {
    "NUMBER": ("return_value_from_key", "value", str),
    "VARIABLE": ("return_value_from_key", "name", None),
    "MULTIPLICATION": ("infix", " * "),
    "DIVISION": ("infix", " / "),
    "PAREN": ("call_function", parse_parenthesis),
    "SUBTRACTION": ("infix", " - "),
    "ADDITION": ("infix", " + "),
    "FUNCTION": ("call_function", parse_function),
}


def apply_infix_operation(node: Node, operation: str) -> str:
    left = parse_ast_node_to_formula(node["left"])
    right = parse_ast_node_to_formula(node["right"])
    return f"{left}{operation}{right}"


def parse_ast_node_to_formula(node: Node) -> str:
    assert isinstance(node, dict), f"Expected a dict, got a {node!r} ({type(node).__name__})"
    if node["type"] not in SERIALIZATION_FUNCTIONS:
        raise NotImplementedError("Unimplemented type {!r}".format(node["type"]))
    serialization_defintion = SERIALIZATION_FUNCTIONS[node["type"]]

    if isinstance(serialization_defintion, tuple):
        # Okay, so it's a structured definition
        command_type, *rest = serialization_defintion
        if command_type == "infix":
            operator_string, *_ = rest
            return apply_infix_operation(node, operator_string)
        elif command_type == "return_value_from_key":
            key_name, cast_function = rest
            value = node[key_name]
            if cast_function is not None:
                value = cast_function(value)
            assert isinstance(
                value, str
            ), f"Tried {key_name!r} in {node}, expected a str but got a {value!r} ({type(value).__name__})"
            return value
        elif command_type == "call_function":
            function, = rest
            value = function(node)
            assert isinstance(
                value, str
            ), f"Tried {function}({node}), expected a str but got a {value!r} ({type(value).__name__})"
            return value
        else:
            raise TypeError("Unrecognized command", command_type)
    elif callable(serialization_defintion):
        value = serialization_defintion(node)
        assert isinstance(
            value, str
        ), f"Called {serialization_defintion}({node}) expecting a str, but got a {value!r} ({type(value).__name__}"
        return value
    raise TypeError("Unknown serialization definition type {serialization_defintion}")


if __name__ == "__main__":
    print(parse_ast_node_to_formula(SAMPLE_TREE))