gadt.py

# Original is https://gist.github.com/msakai/e39c25930b661e8175759cc15175d952
# Forked at 2020-01-09

# Attempt to mimick following GADT (Generalized algebraic data type) in Python:
#
# data Expr a where
#     Const :: a  -> Expr a
#     Add :: Expr Int -> Expr Int -> Expr Int
#     Eq :: Expr Int -> Expr Int -> Expr Bool
#     IfThenElse :: Expr Bool -> Expr a -> Expr a -> Expr a

from dataclasses import dataclass
from typing import Generic, TypeVar

A = TypeVar('A')

class Expr(Generic[A]):
    pass

N = TypeVar('N', int, float)
B = TypeVar('B', bool, bool)  # BUG!
EN = Expr[N]  # Expr[int] will raise error
EB = Expr[B]  # Expr[bool] will raise error


@dataclass(frozen=True)
class Const(Expr[A]):
    value: A

@dataclass(frozen=True)
class Add(Expr[A]):
    left: EN
    right: EN

@dataclass(frozen=True)
class Eq(EB):
    left: EN
    right: EN

@dataclass(frozen=True)
class IfThenElse(Expr[A]):
    condition: EB
    then_expr: Expr[A]
    else_expr: Expr[A]


def eval(e: Expr[A]) -> A:
    if isinstance(e, Const):
        # reveal_type(e)  # => note: Revealed type is 'gadt.Const[Any]'
        return e.value
    elif isinstance(e, Add):
        # reveal_type(e)  # => Revealed type is 'gadt.Add[Any]'
        return eval(e.left) + eval(e.right)
    elif isinstance(e, Eq):
        # reveal_type(e)  # => Revealed type is 'gadt.Eq[Any]'
        return eval(e.left) == eval(e.right)
    elif isinstance(e, IfThenElse):
        # reveal_type(e)  # => note: Revealed type is 'gadt.IfThenElse[Any]'
        return eval(e.then_expr) if eval(e.condition) else eval(e.else_expr)
    else:
        assert False


example = IfThenElse[str](
    Eq(Add(Const[int](2), Const[int](1)), Const[int](0)),
    Const[str]("foo"),
    Const[str]("bar")
)


# reveal_type(example)
# # =>  note: Revealed type is 'gadt.IfThenElse[builtins.str*]'
print(eval(example))