Short circuit If.

Author: dnwpark

tests/test_type_eval.py

@@ -850,6 +850,11 @@ def test_eval_if():
     d = eval_typing(If[Literal[False], int, str])
     assert d is str
 
+    d = eval_typing(If[True, int, str])
+    assert d is int
+    d = eval_typing(If[False, int, str])
+    assert d is str
+
     d = eval_typing(ShorterTuple[tuple[int], tuple[str, str]])
     assert d == tuple[int]
     d = eval_typing(ShorterTuple[tuple[int, int], tuple[str]])

typemap/type_eval/_eval_operators.py

@@ -24,7 +24,6 @@
     GetAttr,
     GreaterThan,
     GreaterThanOrEqual,
-    If,
     IsSubSimilar,
     IsSubtype,
     Iter,
@@ -531,19 +530,6 @@ def func(lhs, rhs, *, ctx):
 ##################################################################
 
 
-@type_eval.register_evaluator(If)
-def _eval_If(cond, then_branch, else_branch, *, ctx):
-    cond_val = _from_literal(cond, ctx)
-    return (
-        _eval_types(then_branch, ctx)
-        if cond_val
-        else _eval_types(else_branch, ctx)
-    )
-
-
-##################################################################
-
-
 def _string_literal_op(typ, op):
     @_lift_over_unions
     def func(*args, ctx):

typemap/type_eval/_eval_typing.py

@@ -17,6 +17,7 @@
     from typing import Any
 
 from . import _apply_generic
+from typemap.typing import If
 
 
 __all__ = ("eval_typing",)
@@ -341,6 +342,14 @@ 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.
+
+    if obj.__origin__ == If:
+        cond = _eval_types(obj.__args__[0], ctx)
+        if cond is True or cond == typing.Literal[True]:
+            return _eval_types(obj.__args__[1], ctx)
+        else:
+            return _eval_types(obj.__args__[2], ctx)
+
     if func := _eval_funcs.get(obj.__origin__):
         new_args = tuple(_eval_types(arg, ctx) for arg in obj.__args__)
         ret = func(*new_args, ctx=ctx)