Simple expression evaluation using AST

expression_evaluation.py

import ast
import json
import os
from typing import Callable, Mapping

FUNCTIONS: Mapping[str, Callable[..., str]] = {
    "concat": lambda *args: "".join(args),
    "json_object": lambda **kwargs: json.dumps(kwargs),
    "env": lambda name: os.getenv(name, ""),
    "format": lambda fmt, *args, **kwargs: fmt.format(*args, **kwargs),
}


def evaluate_expression(expression: ast.expr) -> str:
    expr_type = type(expression)
    if expr_type is ast.Constant:
        return expression.value
    if expr_type is ast.Call:
        if expression.func.id not in FUNCTIONS:
            raise ValueError(f"Unknown function: {expression.func.id}")
        function = FUNCTIONS[expression.func.id]
        args = [evaluate_expression(arg) for arg in expression.args]
        kwargs = {keyword.arg: evaluate_expression(keyword.value) for keyword in expression.keywords}
        return function(*args, **kwargs)
    raise ValueError(f"Unsupported expression type: {expr_type}")


def evaluate(code: str) -> str:
    module = ast.parse(code)
    if len(module.body) != 1:
        raise ValueError("Only one statement is allowed")
    statement = module.body[0]
    if not isinstance(statement, ast.Expr):
        raise ValueError("Only expressions are allowed")
    return evaluate_expression(statement.value)


print(evaluate("format('Gid={};User={}', env('GID'), env('USER'))"))
print(evaluate("json_object(a=1, b=format('Server={};User={}', env('HOST'), env('USER')))"))