use special forms

Author: dnwpark

tests/test_type_eval.py

@@ -854,6 +854,58 @@ def test_eval_if():
     d = eval_typing(ShorterTuple[tuple[int, int], tuple[str]])
     assert d == tuple[str]
 
+    d = eval_typing(If[True, Literal[True], Literal[False]])
+    assert d == Literal[True]
+    d = eval_typing(If[False, Literal[True], Literal[False]])
+    assert d == Literal[False]
+
+
+class Prefixed:
+    x_a: int
+    x_b: str
+    y_a: int
+    y_b: str
+
+
+type FilterPrefix[T, P: str] = NewProtocol[
+    *[
+        x
+        for x in Iter[Attrs[T]]
+        if Equals[StrSlice[GetName[x], 0, Length[P]], P]
+    ]
+]
+type FilterPrefix2[T, Pint: str, Pstr: str] = NewProtocol[
+    *[
+        x
+        for x in Iter[Attrs[T]]
+        if If[
+            Is[GetType[x], int],
+            Equals[StrSlice[GetName[x], 0, Length[Pint]], Pint],
+            Equals[StrSlice[GetName[x], 0, Length[Pstr]], Pstr],
+        ]
+    ]
+]
+
+
+def test_filter_prefix_1():
+    d = eval_typing(FilterPrefix[Prefixed, Literal["x_"]])
+    fmt = format_helper.format_class(d)
+    assert fmt == textwrap.dedent("""\
+        class FilterPrefix[tests.test_type_eval.Prefixed, typing.Literal['x_']]:
+            x_a: int
+            x_b: str
+        """)
+
+
+def test_filter_prefix_2():
+    d = eval_typing(FilterPrefix2[Prefixed, Literal["x_"], Literal["y_"]])
+    fmt = format_helper.format_class(d)
+    assert fmt == textwrap.dedent("""\
+        class FilterPrefix2[tests.test_type_eval.Prefixed, typing.Literal['x_'], typing.Literal['y_']]:
+            x_a: int
+            y_b: str
+        """)
+
 
 def test_uppercase_never():
     d = eval_typing(Uppercase[Never])

typemap/typing.py

@@ -70,46 +70,6 @@ class Length[S: tuple | str]:
     pass
 
 
-class Equals[L, R]:
-    pass
-
-
-class NotEquals[L, R]:
-    pass
-
-
-class GreaterThan[L, R]:
-    pass
-
-
-class LessThan[L, R]:
-    pass
-
-
-class GreaterThanOrEqual[L, R]:
-    pass
-
-
-class LessThanOrEqual[L, R]:
-    pass
-
-
-class Not[T]:
-    pass
-
-
-class And[L, R]:
-    pass
-
-
-class Or[L, R]:
-    pass
-
-
-class If[Cond, Then, Else]:
-    pass
-
-
 class Uppercase[S: str]:
     pass
 
@@ -164,7 +124,16 @@ class _IsGenericAlias(_GenericAlias, _root=True):  # type: ignore[call-arg]
     def __bool__(self):
         evaluator = special_form_evaluator.get()
         if evaluator:
-            return evaluator(self)
+            result = evaluator(self)
+            if result is True or result is False:
+                return result
+            elif (
+                isinstance(result, typing._GenericAlias)
+                and getattr(result, "__origin__", None) is typing.Literal
+            ):
+                return result.__args__[0]
+            else:
+                raise RuntimeError(f"Expected boolean, got {result}")
         else:
             return False
 
@@ -179,4 +148,54 @@ def IsSubSimilar(self, tps):
     return _IsGenericAlias(self, tps)
 
 
+@_SpecialForm
+def Equals(self, params):
+    return _IsGenericAlias(self, params)
+
+
+@_SpecialForm
+def NotEquals(self, params):
+    return _IsGenericAlias(self, params)
+
+
+@_SpecialForm
+def GreaterThan(self, params):
+    return _IsGenericAlias(self, params)
+
+
+@_SpecialForm
+def LessThan(self, params):
+    return _IsGenericAlias(self, params)
+
+
+@_SpecialForm
+def GreaterThanOrEqual(self, params):
+    return _IsGenericAlias(self, params)
+
+
+@_SpecialForm
+def LessThanOrEqual(self, params):
+    return _IsGenericAlias(self, params)
+
+
+@_SpecialForm
+def Not(self, params):
+    return _IsGenericAlias(self, params)
+
+
+@_SpecialForm
+def And(self, params):
+    return _IsGenericAlias(self, params)
+
+
+@_SpecialForm
+def Or(self, params):
+    return _IsGenericAlias(self, params)
+
+
+@_SpecialForm
+def If(self, params):
+    return _IsGenericAlias(self, params)
+
+
 Is = IsSubSimilar