Add bool_special_form.

Author: dnwpark

tests/test_type_eval.py

@@ -33,6 +33,7 @@
     StrConcat,
     StrSlice,
     Uppercase,
+    bool_special_form,
 )
 
 from . import format_helper
@@ -732,9 +733,19 @@ def test_eval_literal_idempotent_01():
         t = nt
 
 
-type IsNotInt[T] = not Is[T, int]
-type IsNotStr[T] = not Is[T, str]
-type IsNotIntOrStr[T] = IsNotInt[T] and IsNotStr[T]
+@bool_special_form
+class IsNotInt[T]:
+    __expr__ = not Is[T, int]
+
+
+@bool_special_form
+class IsNotStr[T]:
+    __expr__ = not Is[T, str]
+
+
+@bool_special_form
+class IsNotIntOrStr[T]:
+    __expr__ = IsNotInt[T] and IsNotStr[T]
 
 
 type SetOfNotInt[T] = set[T] if IsNotInt[T] else T

typemap/type_eval/_eval_typing.py

@@ -19,7 +19,11 @@
     from typing import Any
 
 from . import _apply_generic
-from ._special_form import _special_form_evaluator
+from ._special_form import (
+    BoolSpecialMetadata,
+    _bool_special_form_registry,
+    _special_form_evaluator,
+)
 
 
 __all__ = ("eval_typing",)
@@ -346,13 +350,56 @@ def _eval_applied_type_alias(obj: types.GenericAlias, ctx: EvalContext):
     return evaled
 
 
+def _eval_bool_special_form(
+    metadata: BoolSpecialMetadata,
+    new_args: tuple[typing.Any, ...],
+    ctx: EvalContext,
+) -> bool:
+    import ast
+
+    original_cls = metadata.cls
+
+    try:
+        namespace = {}
+
+        # Add the class's module
+        if cls_module := sys.modules.get(original_cls.__module__):
+            namespace.update(cls_module.__dict__)
+
+        # Add type parameters with their substituted values
+        type_params = metadata.type_params
+        if type_params and new_args:
+            for param, arg in zip(type_params, new_args, strict=False):
+                namespace[param.__name__] = arg
+
+        expr = compile(
+            ast.Expression(body=metadata.expr_node),  # type: ignore[arg-type]
+            '<bool_expr>',
+            'eval',
+        )
+        bool_expr = eval(expr, namespace)
+
+        # Evaluate the type expression
+        result = _eval_types(bool_expr, ctx)
+
+        return result
+
+    except Exception as e:
+        raise RuntimeError(
+            f"Failed to evaluate special form for {original_cls.__name__}: {e}"
+        ) from e
+
+
 @_eval_types_impl.register
 def _eval_applied_class(obj: typing_GenericAlias, ctx: EvalContext):
     """Eval a typing._GenericAlias -- an applied user-defined class"""
     # generic *classes* are typing._GenericAlias while generic type
     # aliases are types.GenericAlias? Why in the world.
     new_args = tuple(_eval_types(arg, ctx) for arg in typing.get_args(obj))
 
+    if metadata := _bool_special_form_registry.get(obj.__origin__):
+        return _eval_bool_special_form(metadata, new_args, ctx)
+
     if func := _eval_funcs.get(obj.__origin__):
         ret = func(*new_args, ctx=ctx)
         # return _eval_types(ret, ctx)  # ???

typemap/type_eval/_special_form.py

@@ -1,4 +1,6 @@
+import ast
 import contextvars
+import dataclasses
 import typing
 from typing import _GenericAlias  # type: ignore
 
@@ -21,10 +23,80 @@ def __iter__(self):
             return iter(typing.TypeVarTuple("_IterDummy"))
 
 
-class _IsGenericAlias(_GenericAlias, _root=True):  # type: ignore[call-arg]
+class _BoolGenericAlias(_GenericAlias, _root=True):  # type: ignore[call-arg]
     def __bool__(self):
         evaluator = _special_form_evaluator.get()
         if evaluator:
             return evaluator(self)
         else:
             return False
+
+
+_IsGenericAlias = _BoolGenericAlias
+
+
+_bool_special_form_registry: dict[typing.Any, BoolSpecialMetadata] = {}
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class BoolSpecialMetadata:
+    cls: type
+    type_params: tuple[type]
+    expr_node: ast.AST
+
+
+def _register_bool_special_form(cls):
+    import inspect
+    import textwrap
+
+    type_params = getattr(cls, '__type_params__', ())
+
+    if '__expr__' not in cls.__dict__:
+        raise TypeError(f"{cls.__name__} must have an '__expr__' field")
+
+    # Parse __expr__ to get the assigned expression
+    source = inspect.getsource(cls)
+    source = textwrap.dedent(source)
+    tree = ast.parse(source)
+
+    expr_node = None
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ClassDef):
+            for item in node.body:
+                if isinstance(item, ast.AnnAssign):
+                    # __expr__: SomeType = expression
+                    if (
+                        isinstance(item.target, ast.Name)
+                        and item.target.id == '__expr__'
+                    ):
+                        expr_node = item.value
+                        break
+                elif isinstance(item, ast.Assign):
+                    # __expr__ = expression
+                    for target in item.targets:
+                        if (
+                            isinstance(target, ast.Name)
+                            and target.id == '__expr__'
+                        ):
+                            expr_node = item.value
+                            break
+                    if expr_node:
+                        break
+            if expr_node:
+                break
+
+    if expr_node is None:
+        raise TypeError(f"Could not find __expr__ assignment in {cls.__name__}")
+
+    def impl_func(self, params):
+        return _BoolGenericAlias(self, params)
+
+    sf = _SpecialForm(impl_func)
+
+    _bool_special_form_registry[sf] = BoolSpecialMetadata(
+        cls=cls,
+        type_params=type_params,
+        expr_node=expr_node,
+    )
+
+    return sf

typemap/typing.py

@@ -4,6 +4,7 @@
     _IterGenericAlias,
     _IsGenericAlias,
     _SpecialForm,
+    _register_bool_special_form,
 )
 
 # Not type-level computation but related
@@ -119,3 +120,7 @@ def IsSubSimilar(self, tps):
 
 
 Is = IsSubSimilar
+
+
+def bool_special_form(cls):
+    return _register_bool_special_form(cls)