spec_parser.py

import ast
import operator
import re
from functools import reduce
from inspect import isawaitable
from typing import Callable
from typing import Dict

replacements = {"!": "not ", "^": " and ", "v": " or "}

pattern = re.compile(r"\!(?!=)|\^|\bv\b")

comparison_repr = {
    operator.eq: "==",
    operator.ne: "!=",
    operator.gt: ">",
    operator.ge: ">=",
    operator.lt: "<",
    operator.le: "<=",
}


def _unique_key(left, right, operator) -> str:
    left_key = getattr(left, "unique_key", "")
    right_key = getattr(right, "unique_key", "")
    return f"{left_key} {operator} {right_key}"


def replace_operators(expr: str) -> str:
    # preprocess the expression adding support for classical logical operators
    def match_func(match):
        return replacements[match.group(0)]

    return pattern.sub(match_func, expr)


def custom_not(predicate: Callable) -> Callable:
    def decorated(*args, **kwargs):
        result = predicate(*args, **kwargs)
        if isawaitable(result):

            async def _negate():
                return not await result

            return _negate()
        return not result

    decorated.__name__ = f"not({predicate.__name__})"
    unique_key = getattr(predicate, "unique_key", "")
    decorated.unique_key = f"not({unique_key})"  # type: ignore[attr-defined]
    return decorated


def custom_and(left: Callable, right: Callable) -> Callable:
    def decorated(*args, **kwargs):
        left_result = left(*args, **kwargs)
        if isawaitable(left_result):

            async def _async_and():
                lr = await left_result
                if not lr:
                    return lr
                rr = right(*args, **kwargs)
                if isawaitable(rr):
                    return await rr
                return rr

            return _async_and()
        if not left_result:
            return left_result
        right_result = right(*args, **kwargs)
        if isawaitable(right_result):
            return right_result
        return right_result

    decorated.__name__ = f"({left.__name__} and {right.__name__})"
    decorated.unique_key = _unique_key(left, right, "and")  # type: ignore[attr-defined]
    return decorated


def custom_or(left: Callable, right: Callable) -> Callable:
    def decorated(*args, **kwargs):
        left_result = left(*args, **kwargs)
        if isawaitable(left_result):

            async def _async_or():
                lr = await left_result
                if lr:
                    return lr
                rr = right(*args, **kwargs)
                if isawaitable(rr):
                    return await rr
                return rr

            return _async_or()
        if left_result:
            return left_result
        right_result = right(*args, **kwargs)
        if isawaitable(right_result):
            return right_result
        return right_result

    decorated.__name__ = f"({left.__name__} or {right.__name__})"
    decorated.unique_key = _unique_key(left, right, "or")  # type: ignore[attr-defined]
    return decorated


def build_constant(constant) -> Callable:
    def decorated(*args, **kwargs):
        return constant

    decorated.__name__ = str(constant)
    decorated.unique_key = str(constant)  # type: ignore[attr-defined]
    return decorated


class Functions:
    registry: Dict[str, Callable] = {}

    @classmethod
    def register(cls, id) -> Callable:
        def register(func):
            cls.registry[id] = func
            return func

        return register

    @classmethod
    def get(cls, func_id):
        func_id = func_id.lower()
        if func_id not in cls.registry:
            raise ValueError(f"Unsupported function: {func_id}")
        return cls.registry[func_id]


class InState:
    def __init__(self, machine):
        self.machine = machine

    def __call__(self, *state_ids: str):
        return set(state_ids).issubset({s.id for s in self.machine.configuration})


@Functions.register("in")
def build_in_call(*state_ids: str) -> Callable:
    state_ids_set = set(state_ids)

    def decorated(*args, **kwargs):
        machine = kwargs["machine"]
        return InState(machine)(*state_ids)

    decorated.__name__ = f"in({state_ids_set})"
    decorated.unique_key = f"in({state_ids_set})"  # type: ignore[attr-defined]
    return decorated


def build_custom_operator(operator) -> Callable:
    operator_repr = comparison_repr[operator]

    def custom_comparator(left: Callable, right: Callable) -> Callable:
        def decorated(*args, **kwargs):
            left_result = left(*args, **kwargs)
            right_result = right(*args, **kwargs)
            if isawaitable(left_result) or isawaitable(right_result):

                async def _async_compare():
                    lr = (await left_result) if isawaitable(left_result) else left_result
                    rr = (await right_result) if isawaitable(right_result) else right_result
                    return bool(operator(lr, rr))

                return _async_compare()
            return bool(operator(left_result, right_result))

        decorated.__name__ = f"({left.__name__} {operator_repr} {right.__name__})"
        decorated.unique_key = _unique_key(left, right, operator_repr)  # type: ignore[attr-defined]
        return decorated

    return custom_comparator


def build_expression(node, variable_hook, operator_mapping):  # noqa: C901
    if isinstance(node, ast.BoolOp):
        # Handle `and` / `or` operations
        operator_fn = operator_mapping[type(node.op)]
        left_expr = build_expression(node.values[0], variable_hook, operator_mapping)
        for right in node.values[1:]:
            right_expr = build_expression(right, variable_hook, operator_mapping)
            left_expr = operator_fn(left_expr, right_expr)
        return left_expr
    elif isinstance(node, ast.Compare):
        # Handle `==` / `!=` / `>` / `<` / `>=` / `<=` operations
        expressions = []
        left_expr = build_expression(node.left, variable_hook, operator_mapping)
        for right_op, right in zip(node.ops, node.comparators):  # noqa: B905  # strict=True requires 3.10+
            right_expr = build_expression(right, variable_hook, operator_mapping)
            operator_fn = operator_mapping[type(right_op)]
            expression = operator_fn(left_expr, right_expr)
            left_expr = right_expr
            expressions.append(expression)

        return reduce(custom_and, expressions)
    elif isinstance(node, ast.Call):
        # Handle function calls
        assert isinstance(node.func, ast.Name)
        constructor = Functions.get(node.func.id)
        params = [arg.value for arg in node.args if isinstance(arg, ast.Constant)]
        return constructor(*params)
    elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not):
        # Handle `not` operation
        operand_expr = build_expression(node.operand, variable_hook, operator_mapping)
        return operator_mapping[type(node.op)](operand_expr)
    elif isinstance(node, ast.Name):
        # Handle variables by calling the variable_hook
        return variable_hook(node.id)
    elif isinstance(node, ast.Constant):
        # Handle constants by returning the value
        return build_constant(node.value)
    else:
        raise ValueError(f"Unsupported expression structure: {node.__class__.__name__}")


def parse_boolean_expr(expr, variable_hook, operator_mapping):
    """Parses the expression into an AST and build a custom expression tree"""
    if expr.strip() == "":
        raise SyntaxError("Empty expression")

    # Optimization trying to avoid parsing the expression if not needed
    if "!" not in expr and " " not in expr and "In(" not in expr:
        return variable_hook(expr)
    expr = replace_operators(expr)
    tree = ast.parse(expr, mode="eval")
    return build_expression(tree.body, variable_hook, operator_mapping)


operator_mapping = {
    ast.Or: custom_or,
    ast.And: custom_and,
    ast.Not: custom_not,
    ast.GtE: build_custom_operator(operator.ge),
    ast.Gt: build_custom_operator(operator.gt),
    ast.LtE: build_custom_operator(operator.le),
    ast.Lt: build_custom_operator(operator.lt),
    ast.Eq: build_custom_operator(operator.eq),
    ast.NotEq: build_custom_operator(operator.ne),
}