Mini language for simple operations

minilang.py

import operator
from functools import partialmethod


class Forward:

    def resolve(self, env):
        raise NotImplementedError()

    def op(self, op, rhs):
        if not isinstance(rhs, Forward):
            raise ValueError('"{}" not of type `Forward`'.format(rhs))
        return Evaluation(op, self, rhs)

    __eq__ = partialmethod(op, operator.__eq__)
    __gt__ = partialmethod(op, operator.__gt__)
    __lt__ = partialmethod(op, operator.__lt__)
    __ge__ = partialmethod(op, operator.__ge__)
    __le__ = partialmethod(op, operator.__le__)

    __add__ = partialmethod(op, operator.__add__)
    __sub__ = partialmethod(op, operator.__sub__)
    __mul__ = partialmethod(op, operator.__mul__)
    __truediv__ = partialmethod(op, operator.__truediv__)
    __floordiv__ = partialmethod(op, operator.__floordiv__)
    __pow__ = partialmethod(op, operator.__pow__)

    __iadd__ = partialmethod(op, operator.__iadd__)
    __isub__ = partialmethod(op, operator.__isub__)
    __imul__ = partialmethod(op, operator.__imul__)
    __itruediv__ = partialmethod(op, operator.__itruediv__)
    __ifloordiv__ = partialmethod(op, operator.__ifloordiv__)
    __ipow__ = partialmethod(op, operator.__ipow__)


class Evaluation(Forward):

    def __init__(self, op, lhs, rhs):
        self.op = op
        self.lhs = lhs
        self.rhs = rhs

    def resolve(self, env):
        return self.op(self.lhs.resolve(env), self.rhs.resolve(env))


class If(Forward):

    def __init__(self, evaluation, then, else_):
        self.evaluation = evaluation
        self.then = then
        self.else_ = else_
        
    def resolve(self, env):
        if self.evaluation.resolve(env):
            return self.then.resolve(env)
        return self.else_.resolve(env)


class Constant(Forward):
    def __init__(self, value):
        self.value = value

    def resolve(self, env):
        return self.value


class Variable(Forward):

    def __init__(self, name):
        self.name = name

    def resolve(self, env):
        return env[self.name]


class Optional(Variable):

    def resolve(self, env):
        return env.get(self.name)


def resolve(evaluation, env):
    return evaluation.resolve(env)


TESTS = [
    ("Constant(5) + Constant(10)", {}, 15),
    ("Variable('foo') + Constant(10)", {'foo': 3}, 13),
    ("Constant(5) + Constant(10) + Constant(5)", {}, 20),
    ("Constant(5) + Constant(10) * Constant(5)", {}, 55),
    ("(Constant(5) + Constant(10)) * Constant(5)", {}, 75),
    (
        "(Variable('foo') + Variable('bar')) * Constant(5)",
        {'foo': 2, 'bar': 3},
        25
    ),
    ("Variable('foo') == Variable('bar')", {'foo': 2, 'bar': 3}, False),
    (
        """
        If(
            Variable('foo') == Variable('bar'),
            Constant(True),
            Constant(False))
        """,
        {'foo': 2, 'bar': 3},
        False
    ),
    (
        """
        If(
            Optional('bar'),
            Variable('bar'),
            Variable('foo'))
        """,
        {'foo': 20},
        20
    ),
    (
        """
        If(
            Optional('bar'),
            Variable('bar'),
            Variable('foo'))
        """,
        {'foo': 20, 'bar': 10},
        10
    ),
    (
        """
        Constant(5) + If(Optional('foo'), Variable('foo'), Constant(10))
        """,
        {},
        15
    )
]

if __name__ == "__main__":
    for func, env, result in TESTS:
        print(func, ' --> ', env, ' ->> ', result)
        try:
            assert (
                eval("resolve({}, env)".format(func), globals(), locals())
                == result)
        except AssertionError:
            print('Error')
            print('"{}" != "{}"'.format(
                eval("resolve({}, env)".format(func), globals(), locals()),
                result
            ))