Add test_astlike_1.

Author: dnwpark

tests/test_astlike_1.py

@@ -0,0 +1,346 @@
+import typing
+
+from typemap.type_eval import eval_call_with_types, eval_typing
+from typemap.typing import (
+    Attrs,
+    BaseTypedDict,
+    Bool,
+    GetArg,
+    GetName,
+    GetType,
+    IsSub,
+    Iter,
+    Matches,
+    Member,
+    NewProtocol,
+    _BoolLiteral,
+)
+
+
+"""
+An AST like system for doing simple type checked computation.
+
+Provides Constant and Variable nodes which can be used with operators
+(+, -, *, /, //, **, and %) to build up expression trees.
+
+Calling eval on a Node with **kwargs corresponding to the Variables in
+the expression will compute and return the result.
+
+Example usage:
+    a = Variable[int, typing.Literal["a"]]()
+    b = Variable[int, typing.Literal["b"]]()
+    c = Constant(3)
+    z = a + b * c
+    result = eval(z, a=1, b=2)
+    assert result is 7
+"""
+
+
+type VarArg[Name: str, T: type] = tuple[Name, T]
+
+type VarArgName[V: VarArg] = GetArg[V, tuple, typing.Literal[0]]
+type VarArgType[V: VarArg] = GetArg[V, tuple, typing.Literal[1]]
+
+
+type CombineVarArgs[Ls: tuple[VarArg], Rs: tuple[VarArg]] = tuple[
+    *[
+        VarArg[
+            VarArgName[x],
+            (
+                VarArgType[x]
+                if not any(  # Unique to Ls
+                    Matches[VarArgName[x], VarArgName[y]] for y in Iter[Rs]
+                )
+                else GetArg[  # Common to both Ls and Rs
+                    tuple[
+                        *[
+                            (
+                                VarArgType[x]
+                                if IsSub[VarArgType[x], VarArgType[y]]
+                                else VarArgType[y]
+                                if IsSub[VarArgType[y], VarArgType[x]]
+                                else typing.Never
+                            )
+                            for y in Iter[Rs]
+                            if Matches[VarArgName[x], VarArgName[y]]
+                        ]
+                    ],
+                    tuple,
+                    typing.Literal[0],
+                ]
+            ),
+        ]
+        for x in Iter[Ls]
+    ],
+    *[  # Unique to Rs
+        x
+        for x in Iter[Rs]
+        if not any(  # Unique to Rs
+            Matches[VarArgName[x], VarArgName[y]] for y in Iter[Ls]
+        )
+    ],
+]
+
+
+def test_astlike_1_combine_varargs_01():
+    t = eval_typing(
+        CombineVarArgs[
+            tuple[
+                VarArg[typing.Literal["same"], int],
+                VarArg[typing.Literal["different"], int],
+                VarArg[typing.Literal["left_sub"], bool],
+                VarArg[typing.Literal["right_sub"], int],
+                VarArg[typing.Literal["unique_left"], int],
+            ],
+            tuple[
+                VarArg[typing.Literal["same"], int],
+                VarArg[typing.Literal["different"], float],
+                VarArg[typing.Literal["left_sub"], int],
+                VarArg[typing.Literal["right_sub"], bool],
+                VarArg[typing.Literal["unique_right"], int],
+            ],
+        ]
+    )
+    assert (
+        t
+        == tuple[
+            tuple[typing.Literal["same"], int],
+            tuple[typing.Literal["different"], typing.Never],
+            tuple[typing.Literal["left_sub"], bool],
+            tuple[typing.Literal["right_sub"], bool],
+            tuple[typing.Literal["unique_left"], int],
+            tuple[typing.Literal["unique_right"], int],
+        ]
+    )
+
+
+type IsAssignable[L, R] = (
+    IsSub[R, L]
+    or Bool[Matches[L, float] and Bool[IsFloat[R]]]
+    or Bool[Matches[L, complex] and Bool[IsComplex[R]]]
+)
+type VarIsPresent[V: VarArg, K: BaseTypedDict] = any(
+    Matches[VarArgName[V], GetName[x]]
+    and Bool[IsAssignable[VarArgType[V], GetType[x]]]
+    for x in Iter[Attrs[K]]
+)
+type AllVarsPresent[Vs: tuple[VarArg, ...], K: BaseTypedDict] = all(
+    Bool[VarIsPresent[v, K]] for v in Iter[Vs]
+)
+
+
+def test_astlike_1_all_vars_present_01():
+    t = eval_typing(
+        AllVarsPresent[
+            tuple[VarArg[typing.Literal["x"], int]],
+            NewProtocol[Member[typing.Literal["x"], int]],
+        ]
+    )
+    assert t == _BoolLiteral[True]
+
+
+def test_astlike_1_all_vars_present_02():
+    t = eval_typing(
+        AllVarsPresent[
+            tuple[VarArg[typing.Literal["x"], int]],
+            NewProtocol[Member[typing.Literal["x"], bool]],
+        ]
+    )
+    assert t == _BoolLiteral[True]
+
+
+def test_astlike_1_all_vars_present_03():
+    t = eval_typing(
+        AllVarsPresent[
+            tuple[VarArg[typing.Literal["x"], int]],
+            NewProtocol[Member[typing.Literal["x"], float]],
+        ]
+    )
+    assert t == _BoolLiteral[False]
+
+
+def test_astlike_1_all_vars_present_04():
+    t = eval_typing(
+        AllVarsPresent[
+            tuple[VarArg[typing.Literal["x"], int]],
+            NewProtocol[
+                Member[typing.Literal["x"], int],
+                Member[typing.Literal["y"], int],
+            ],
+        ]
+    )
+    assert t == _BoolLiteral[True]
+
+
+def test_astlike_1_all_vars_present_05():
+    t = eval_typing(
+        AllVarsPresent[
+            tuple[VarArg[typing.Literal["x"], int]],
+            NewProtocol[Member[typing.Literal["y"], int]],
+        ]
+    )
+    assert t == _BoolLiteral[False]
+
+
+type IsIntegral[T] = IsSub[T, int]
+type IsFloat[T] = Bool[IsIntegral[T]] or IsSub[T, float]
+type IsComplex[T] = Bool[IsFloat[T]] or IsSub[T, complex]
+
+type SimpleNumericOp[L, R] = (
+    int
+    if Bool[IsIntegral[L]] and Bool[IsIntegral[R]]
+    else float
+    if Bool[IsFloat[L]] and Bool[IsFloat[R]]
+    else typing.Never
+)
+type ComplexNumericOp[L, R] = (
+    SimpleNumericOp[L, R]
+    if Bool[IsFloat[L]] and Bool[IsFloat[R]]
+    else complex
+    if Bool[IsComplex[L]] and Bool[IsComplex[R]]
+    else typing.Never
+)
+
+type Add[L, R] = (
+    ComplexNumericOp[L, R]
+    if Bool[IsComplex[L]] and Bool[IsComplex[R]]
+    else typing.Never
+)
+type Sub[L, R] = (
+    ComplexNumericOp[L, R]
+    if Bool[IsComplex[L]] and Bool[IsComplex[R]]
+    else typing.Never
+)
+type Mul[L, R] = (
+    ComplexNumericOp[L, R]
+    if Bool[IsComplex[L]] and Bool[IsComplex[R]]
+    else typing.Never
+)
+type TrueDiv[L, R] = (
+    float
+    if IsSub[L, int] and IsSub[R, int]
+    else ComplexNumericOp[L, R]
+    if Bool[IsComplex[L]] and Bool[IsComplex[R]]
+    else typing.Never
+)
+type FloorDiv[L, R] = ComplexNumericOp[L, R]
+type Pow[L, R] = ComplexNumericOp[L, R]
+type Mod[L, R] = (
+    SimpleNumericOp[L, R]
+    if Bool[IsFloat[L]] and Bool[IsFloat[R]]
+    else typing.Never
+)
+
+
+class NodeMeta(type): ...
+
+
+class Node[T, Vs: tuple[VarArg, ...]](metaclass=NodeMeta):
+    def __add__[OtherT, OtherVs: tuple[VarArg, ...]](
+        self, other: Node[OtherT, OtherVs]
+    ) -> Node[Add[T, OtherT], CombineVarArgs[Vs, OtherVs]]: ...
+
+    def __sub__[OtherT, OtherVs: tuple[VarArg, ...]](
+        self, other: Node[OtherT, OtherVs]
+    ) -> Node[Sub[T, OtherT], CombineVarArgs[Vs, OtherVs]]: ...
+
+    def __mul__[OtherT, OtherVs: tuple[VarArg, ...]](
+        self, other: Node[OtherT, OtherVs]
+    ) -> Node[Mul[T, OtherT], CombineVarArgs[Vs, OtherVs]]: ...
+
+    def __truediv__[OtherT, OtherVs: tuple[VarArg, ...]](
+        self, other: Node[OtherT, OtherVs]
+    ) -> Node[TrueDiv[T, OtherT], CombineVarArgs[Vs, OtherVs]]: ...
+
+    def __floordiv__[OtherT, OtherVs: tuple[VarArg, ...]](
+        self, other: Node[OtherT, OtherVs]
+    ) -> Node[FloorDiv[T, OtherT], CombineVarArgs[Vs, OtherVs]]: ...
+
+    def __pow__[OtherT, OtherVs: tuple[VarArg, ...]](
+        self, other: Node[OtherT, OtherVs]
+    ) -> Node[Pow[T, OtherT], CombineVarArgs[Vs, OtherVs]]: ...
+
+    def __mod__[OtherT, OtherVs: tuple[VarArg, ...]](
+        self, other: Node[OtherT, OtherVs]
+    ) -> Node[Mod[T, OtherT], CombineVarArgs[Vs, OtherVs]]: ...
+
+
+class Constant[T](Node[T, tuple[()]]):
+    value: typing.Any
+
+    def __init__(self, value) -> None:
+        self.value = value
+
+
+def test_astlike_1_constant_01():
+    t = eval_typing(Constant[int])
+    assert t == Constant[int]
+
+
+def test_astlike_1_constant_02():
+    t = eval_call_with_types(eval, Constant[int])
+    assert t is int
+
+    t = eval_call_with_types(eval, Constant[int], x=1)
+    assert t is int
+
+
+class Variable[T, Name: typing.Literal[str]](Node[T, tuple[VarArg[Name, T]]]):
+    @property
+    def name(self) -> typing.Literal[Name]:
+        return self.__orig_class__.__args__[1].__args__[0]
+
+    def _eval(self, **kwargs) -> T:
+        if self.name not in kwargs:
+            raise ValueError(f"Expected '{self.name}' in kwargs")
+        if not isinstance(kwargs[self.name], self.__orig_class__.__args__[0]):
+            raise ValueError(
+                f"Expected '{self.__orig_class__.__args__[0].__name__}', "
+                f"got '{type(kwargs[self.name]).__name__}'"
+            )
+        return kwargs[self.name]
+
+
+def test_astlike_1_variable_01():
+    n = Variable[int, typing.Literal["x"]]
+    assert n().name == "x"
+
+
+def test_astlike_1_variable_02():
+    t = eval_call_with_types(eval, Variable[int, typing.Literal["x"]], x=int)
+    assert t is int
+    t = eval_call_with_types(eval, Variable[int, typing.Literal["x"]], x=bool)
+    assert t is int
+    t = eval_call_with_types(eval, Variable[int, typing.Literal["x"]], x=str)
+    assert t is typing.Never
+
+
+def eval[T, Vs: tuple[VarArg, ...], K: BaseTypedDict](
+    self: Node[T, Vs], **kwargs: typing.Unpack[K]
+) -> T if Bool[AllVarsPresent[Vs, K]] else typing.Never: ...
+
+
+def test_astlike_1_eval_01():
+    n = Node[int, tuple[VarArg[typing.Literal["x"], int]]]
+    t = eval_call_with_types(eval, n, x=int)
+    assert t is int
+    t = eval_call_with_types(eval, n, x=bool)
+    assert t is int
+    t = eval_call_with_types(eval, n, x=str)
+    assert t is typing.Never
+
+
+def test_astlike_1_eval_02():
+    n = Node[
+        complex,
+        tuple[
+            VarArg[typing.Literal["x"], float],
+            VarArg[typing.Literal["y"], float],
+        ],
+    ]
+    t = eval_call_with_types(eval, n, x=int, y=int)
+    assert t is complex
+    t = eval_call_with_types(eval, n, x=bool, y=float)
+    assert t is complex
+    t = eval_call_with_types(eval, n, x=str, y=complex)
+    assert t is typing.Never