Expression Evaluator

expr_eval.py

import enum
import contextvars
import contextlib
import operator as ops
from dataclasses import dataclass
from pyparsing import (
    Word,
    nums,
    oneOf,
    infixNotation,
    opAssoc,
    Literal,
    Forward,
    QuotedString,
    Regex,
    ParseResults
)


class TokenType(str, enum.Enum):
    INTEGER = "integer"
    FLOAT = "float"
    STRING = "string"
    BINARY_OPERATOR = "binary_operator"
    VARIABLE = "variable"

OPERATOR_IMPLS =  {
    "+": ops.add,
    "-": ops.sub,
    "/": ops.truediv,
    "*": ops.mul,
    "AND": lambda a, b: a and b,
    "OR": lambda a, b: a or b,
    "==": ops.eq,
    ">": ops.gt,
    "<": ops.lt,
    ">=": ops.ge,
    "<=": ops.le,
}

OPERATORS = OPERATOR_IMPLS.keys()

# index represents the operator precedence
# 0th index has least precedence
OPERATOR_PRECEDENCES = [
    ["AND", "OR"],
    ["==", ">", "<", ">=", "<="],
    ["+", "-"],
    ["/", "*"]
]

TokenCasters = {
    TokenType.INTEGER: int,
    TokenType.FLOAT: float,
    TokenType.STRING: str,
    TokenType.BINARY_OPERATOR: lambda op: OPERATOR_IMPLS[op],
    TokenType.VARIABLE: lambda var: variable_lookup_context.get()[var],
}


@dataclass
class TokenizationResult:
    type: TokenType
    value: str

    precedence: int

    @property
    def parsed_value(self):
        return TokenCasters[self.type](self.value)


def tokenization_result(type: TokenType):
    def inner(t):
        precedence = -1
        for prec_value, symbols in enumerate(OPERATOR_PRECEDENCES):
            if t[0] in symbols:
                precedence = prec_value
                break

        return TokenizationResult(type=type, value=t[0], precedence=precedence)

    return inner


floats = Regex(r"\d+\.\d*|\d*\.?\d+").setParseAction(
    tokenization_result(TokenType.FLOAT)
)
ints = Word(nums).setParseAction(tokenization_result(TokenType.INTEGER))
binary_operator = oneOf(OPERATORS).setParseAction(
    tokenization_result(TokenType.BINARY_OPERATOR)
)

left_paren = Literal("(").suppress()
right_paren = Literal(")").suppress()

variable_string = QuotedString(quoteChar='$"', endQuoteChar='"').setParseAction(
    tokenization_result(TokenType.VARIABLE)
)
quoted_string = QuotedString(quoteChar='"', endQuoteChar='"').setParseAction(
    tokenization_result(TokenType.STRING)
)

expr = Forward()
expr <<= infixNotation(
    variable_string  | ints | floats | quoted_string,
    [
        (binary_operator, 2, opAssoc.LEFT)
    ],
)

variable_lookup_context = contextvars.ContextVar("variable_lookup", default=None)

def _infix_evaluator(eval_string: str):
    def _helper(tokens):
        operands = []
        operators = []
        for token in tokens:
            if isinstance(token, ParseResults):
                operands.append(_helper(token))
            elif token.type is TokenType.BINARY_OPERATOR:
                while operators and operators[-1].precedence >= token.precedence:
                    fn = operators.pop().parsed_value
                    right_opn = operands.pop()
                    left_opn = operands.pop()
                    operands.append(fn(left_opn, right_opn))
                operators.append(token)
            else:
                operands.append(token.parsed_value)
            
        while operators:
            fn = operators.pop().parsed_value
            right_opn = operands.pop()
            left_opn = operands.pop()
            
            operands.append(fn(left_opn, right_opn))

        return operands[0]

    tokens = expr.parseString(eval_string)
    return _helper(tokens)

@contextlib.contextmanager
def evaluation_ctx(variables: dict):
    try:
        variable_lookup_context.set(variables)
        yield _infix_evaluator
    finally:
        variable_lookup_context.set(None)

input_string = '1 + 2 <= 3 AND (4 > 10 OR 100 < $"foo") AND ($"bar" == "abcd")'

with evaluation_ctx({"foo": 1234, "bar": "abcd"}) as evaluator:
    print(evaluator(input_string))