refactor visit_conditional_expr

Author: randolf-scholz

mypy/checkexpr.py

@@ -141,7 +141,6 @@
     get_type_vars,
     is_literal_type_like,
     make_simplified_union,
-    simple_literal_type,
     true_only,
     try_expanding_sum_type_to_union,
     try_getting_str_literals,
@@ -5888,63 +5887,26 @@ def visit_conditional_expr(self, e: ConditionalExpr, allow_none_return: bool = F
             elif else_map is None:
                 self.msg.redundant_condition_in_if(True, e.cond)
 
+        if ctx is None:
+            # When no context is provided, compute each branch individually, and
+            # use the union of the results as artificial context. Important for:
+            # - testUnificationDict
+            # - testConditionalExpressionWithEmpty
+            ctx_if_type = self.analyze_cond_branch(
+                if_map, e.if_expr, context=ctx, allow_none_return=allow_none_return
+            )
+            ctx_else_type = self.analyze_cond_branch(
+                else_map, e.else_expr, context=ctx, allow_none_return=allow_none_return
+            )
+            ctx = make_simplified_union([ctx_if_type, ctx_else_type])
+
         if_type = self.analyze_cond_branch(
             if_map, e.if_expr, context=ctx, allow_none_return=allow_none_return
         )
-
-        # we want to keep the narrowest value of if_type for union'ing the branches
-        # however, it would be silly to pass a literal as a type context. Pass the
-        # underlying fallback type instead.
-        if_type_fallback = simple_literal_type(get_proper_type(if_type)) or if_type
-
-        # Analyze the right branch using full type context and store the type
-        full_context_else_type = self.analyze_cond_branch(
+        else_type = self.analyze_cond_branch(
             else_map, e.else_expr, context=ctx, allow_none_return=allow_none_return
         )
 
-        if not mypy.checker.is_valid_inferred_type(if_type, self.chk.options):
-            # Analyze the right branch disregarding the left branch.
-            else_type = full_context_else_type
-            # we want to keep the narrowest value of else_type for union'ing the branches
-            # however, it would be silly to pass a literal as a type context. Pass the
-            # underlying fallback type instead.
-            else_type_fallback = simple_literal_type(get_proper_type(else_type)) or else_type
-
-            # If it would make a difference, re-analyze the left
-            # branch using the right branch's type as context.
-            if ctx is None or not is_equivalent(else_type_fallback, ctx):
-                # TODO: If it's possible that the previous analysis of
-                # the left branch produced errors that are avoided
-                # using this context, suppress those errors.
-                if_type = self.analyze_cond_branch(
-                    if_map,
-                    e.if_expr,
-                    context=else_type_fallback,
-                    allow_none_return=allow_none_return,
-                )
-
-        elif if_type_fallback == ctx:
-            # There is no point re-running the analysis if if_type is equal to ctx.
-            # That would  be an exact duplicate of the work we just did.
-            # This optimization is particularly important to avoid exponential blowup with nested
-            # if/else expressions: https://github.com/python/mypy/issues/9591
-            # TODO: would checking for is_proper_subtype also work and cover more cases?
-            else_type = full_context_else_type
-        else:
-            # Analyze the right branch in the context of the left
-            # branch's type.
-            else_type = self.analyze_cond_branch(
-                else_map,
-                e.else_expr,
-                context=if_type_fallback,
-                allow_none_return=allow_none_return,
-            )
-
-        # In most cases using if_type as a context for right branch gives better inferred types.
-        # This is however not the case for literal types, so use the full context instead.
-        if is_literal_type_like(full_context_else_type) and not is_literal_type_like(else_type):
-            else_type = full_context_else_type
-
         res: Type = make_simplified_union([if_type, else_type])
         if has_uninhabited_component(res) and not isinstance(
             get_proper_type(self.type_context[-1]), UnionType

mypyc/irbuild/statement.py

@@ -599,7 +599,9 @@ def make_entry(type: Expression) -> tuple[ValueGenFunc, int]:
         (make_entry(type) if type else None, var, make_handler(body))
         for type, var, body in zip(t.types, t.vars, t.handlers)
     ]
-    else_body = (lambda: builder.accept(t.else_body)) if t.else_body else None
+
+    _else_body = t.else_body
+    else_body = (lambda: builder.accept(_else_body)) if _else_body else None
     transform_try_except(builder, body, handlers, else_body, t.line)
 
 

test-data/unit/check-literal.test

@@ -2946,6 +2946,140 @@ reveal_type(C().collection)  # N: Revealed type is "builtins.list[Literal['word'
 reveal_type(C().word)  # N: Revealed type is "Literal['word']"
 [builtins fixtures/tuple.pyi]
 
+[case testStringLiteralTernary]
+def test(b: bool) -> None:
+    l = "foo" if b else "bar"
+    reveal_type(l)  # N: Revealed type is "builtins.str"
+[builtins fixtures/tuple.pyi]
+
+[case testintLiteralTernary]
+def test(b: bool) -> None:
+    l = 0 if b else 1
+    reveal_type(l)  # N: Revealed type is "builtins.int"
+[builtins fixtures/tuple.pyi]
+
+[case testStringIntUnionTernary]
+def test(b: bool) -> None:
+    l = 1 if b else "a"
+    reveal_type(l)  # N: Revealed type is "Union[builtins.int, builtins.str]"
+[builtins fixtures/tuple.pyi]
+
+[case testListComprehensionTernary]
+# gh-19534
+def test(b: bool) -> None:
+    l = [1] if b else ["a"]
+    reveal_type(l)  # N: Revealed type is "Union[builtins.list[builtins.int], builtins.list[builtins.str]]"
+[builtins fixtures/list.pyi]
+
+[case testSetComprehensionTernary]
+# gh-19534
+def test(b: bool) -> None:
+    s = {1} if b else {"a"}
+    reveal_type(s)  # N: Revealed type is "Union[builtins.set[builtins.int], builtins.set[builtins.str]]"
+[builtins fixtures/set.pyi]
+
+[case testDictComprehensionTernary]
+# gh-19534
+def test(b: bool) -> None:
+    d = {1:1} if "" else {"a": "a"}
+    reveal_type(d)  # N: Revealed type is "Union[builtins.dict[builtins.int, builtins.int], builtins.dict[builtins.str, builtins.str]]"
+[builtins fixtures/dict.pyi]
+
+[case testLambdaTernary]
+from typing import TypeVar, Union, Callable, reveal_type
+
+NOOP = lambda: None
+class A: pass
+class B:
+    attr: Union[A, None]
+
+def test_static(x: Union[A, None]) -> None:
+    def foo(t: A) -> None: ...
+
+    l1: Callable[[], object] = (lambda: foo(x)) if x is not None else NOOP
+    r1: Callable[[], object] = NOOP if x is None else (lambda: foo(x))
+    l2 = (lambda: foo(x)) if x is not None else NOOP
+    r2 = NOOP if x is None else (lambda: foo(x))
+    reveal_type(l2)  # N: Revealed type is "def ()"
+    reveal_type(r2)  # N: Revealed type is "def ()"
+
+def test_generic(x: Union[A, None]) -> None:
+    T = TypeVar("T")
+    def bar(t: T) -> T: return t
+
+    l1: Callable[[], None] = (lambda: bar(x)) if x is None else NOOP
+    r1: Callable[[], None] = NOOP if x is not None else (lambda: bar(x))
+    l2 = (lambda: bar(x)) if x is None else NOOP
+    r2 = NOOP if x is not None else (lambda: bar(x))
+    reveal_type(l2)  # N: Revealed type is "def ()"
+    reveal_type(r2)  # N: Revealed type is "def ()"
+
+
+[case testLambdaTernaryIndirectAttribute]
+# fails due to binder issue inside `check_func_def`
+# gh-19561
+from typing import TypeVar, Union, Callable, reveal_type
+
+NOOP = lambda: None
+class A: pass
+class B:
+    attr: Union[A, None]
+
+def test_static_with_attr(x: B) -> None:
+    def foo(t: A) -> None: ...
+
+    l1: Callable[[], None] = (lambda: foo(x.attr)) if x.attr is not None else NOOP  # E: Argument 1 to "foo" has incompatible type "Optional[A]"; expected "A"
+    r1: Callable[[], None] = NOOP if x.attr is None else (lambda: foo(x.attr))  # E: Argument 1 to "foo" has incompatible type "Optional[A]"; expected "A"
+    l2 = (lambda: foo(x.attr)) if x.attr is not None else NOOP  # E: Argument 1 to "foo" has incompatible type "Optional[A]"; expected "A"
+    r2 = NOOP if x.attr is None else (lambda: foo(x.attr))  # E: Argument 1 to "foo" has incompatible type "Optional[A]"; expected "A"
+    reveal_type(l2)  # N: Revealed type is "def ()"
+    reveal_type(r2)  # N: Revealed type is "def ()"
+
+def test_generic_with_attr(x: B) -> None:
+    T = TypeVar("T")
+    def bar(t: T) -> T: return t
+
+    l1: Callable[[], None] = (lambda: bar(x.attr)) if x.attr is None else NOOP  # E: Incompatible types in assignment (expression has type "Callable[[], Optional[A]]", variable has type "Callable[[], None]")
+    r1: Callable[[], None] = NOOP if x.attr is not None else (lambda: bar(x.attr))  # E: Incompatible types in assignment (expression has type "Callable[[], Optional[A]]", variable has type "Callable[[], None]")
+    l2 = (lambda: bar(x.attr)) if x.attr is None else NOOP
+    r2 = NOOP if x.attr is not None else (lambda: bar(x.attr))
+    reveal_type(l2)  # N: Revealed type is "def () -> Union[__main__.A, None]"
+    reveal_type(r2)  # N: Revealed type is "def () -> Union[__main__.A, None]"
+
+[case testLambdaTernaryDoubleIndirectAttribute]
+# fails due to binder issue inside `check_func_def`
+# gh-19561
+from typing import TypeVar, Union, Callable, reveal_type
+
+NOOP = lambda: None
+class A: pass
+class B:
+    attr: Union[A, None]
+class C:
+    attr: B
+
+def test_static_with_attr(x: C) -> None:
+    def foo(t: A) -> None: ...
+
+    l1: Callable[[], None] = (lambda: foo(x.attr.attr)) if x.attr.attr is not None else NOOP  # E: Argument 1 to "foo" has incompatible type "Optional[A]"; expected "A"
+    r1: Callable[[], None] = NOOP if x.attr.attr is None else (lambda: foo(x.attr.attr))  # E: Argument 1 to "foo" has incompatible type "Optional[A]"; expected "A"
+    l2 = (lambda: foo(x.attr.attr)) if x.attr.attr is not None else NOOP  # E: Argument 1 to "foo" has incompatible type "Optional[A]"; expected "A"
+    r2 = NOOP if x.attr.attr is None else (lambda: foo(x.attr.attr))  # E: Argument 1 to "foo" has incompatible type "Optional[A]"; expected "A"
+    reveal_type(l2)  # N: Revealed type is "def ()"
+    reveal_type(r2)  # N: Revealed type is "def ()"
+
+def test_generic_with_attr(x: C) -> None:
+    T = TypeVar("T")
+    def bar(t: T) -> T: return t
+
+    l1: Callable[[], None] = (lambda: bar(x.attr.attr)) if x.attr.attr is None else NOOP  # E: Incompatible types in assignment (expression has type "Callable[[], Optional[A]]", variable has type "Callable[[], None]")
+    r1: Callable[[], None] = NOOP if x.attr.attr is not None else (lambda: bar(x.attr.attr))  # E: Incompatible types in assignment (expression has type "Callable[[], Optional[A]]", variable has type "Callable[[], None]")
+    l2 = (lambda: bar(x.attr.attr)) if x.attr.attr is None else NOOP
+    r2 = NOOP if x.attr.attr is not None else (lambda: bar(x.attr.attr))
+    reveal_type(l2)  # N: Revealed type is "def () -> Union[__main__.A, None]"
+    reveal_type(r2)  # N: Revealed type is "def () -> Union[__main__.A, None]"
+
+
 [case testLiteralTernaryUnionNarrowing]
 from typing import Literal, Optional
 

test-data/unit/check-optional.test

@@ -427,7 +427,9 @@ reveal_type(l)  # N: Revealed type is "builtins.list[typing.Generator[builtins.s
 [builtins fixtures/list.pyi]
 
 [case testNoneListTernary]
-x = [None] if "" else [1]  # E: List item 0 has incompatible type "int"; expected "None"
+# gh-19534
+x = [None] if "" else [1]
+reveal_type(x)  # N: Revealed type is "Union[builtins.list[None], builtins.list[builtins.int]]"
 [builtins fixtures/list.pyi]
 
 [case testListIncompatibleErrorMessage]