[red-knot] Type narrow in else clause (#13918)

## Summary

Add support for type narrowing in elif and else scopes as part of
#13694.

## Test Plan

- mdtest
- builder unit test for union negation.

---------

Co-authored-by: Carl Meyer <[email protected]>

crates/red_knot_python_semantic/resources/mdtest/comparison/integers.md

@@ -19,7 +19,8 @@ reveal_type(1 <= "" and 0 < 1)  # revealed: @Todo | Literal[True]
 
 ```py
 # TODO: implement lookup of `__eq__` on typeshed `int` stub.
-def int_instance() -> int: ...
+def int_instance() -> int:
+    return 42
 
 reveal_type(1 == int_instance())  # revealed: @Todo
 reveal_type(9 < int_instance())  # revealed: bool

crates/red_knot_python_semantic/resources/mdtest/comparison/tuples.md

@@ -59,7 +59,8 @@ reveal_type(c >= d)  # revealed: Literal[True]
 
 ```py
 def bool_instance() -> bool: ...
-def int_instance() -> int: ...
+def int_instance() -> int:
+    return 42
 
 a = (bool_instance(),)
 b = (int_instance(),)
@@ -159,7 +160,8 @@ reveal_type(a >= a)  # revealed: @Todo
 "Membership Test Comparisons" refers to the operators `in` and `not in`.
 
 ```py
-def int_instance() -> int: ...
+def int_instance() -> int:
+    return 42
 
 a = (1, 2)
 b = ((3, 4), (1, 2))

crates/red_knot_python_semantic/resources/mdtest/narrow/conditionals_elif_else.md

@@ -0,0 +1,57 @@
+# Narrowing for conditionals with elif and else
+
+## Positive contributions become negative in elif-else blocks
+
+```py
+def int_instance() -> int:
+    return 42
+
+x = int_instance()
+
+if x == 1:
+    # cannot narrow; could be a subclass of `int`
+    reveal_type(x)  # revealed: int
+elif x == 2:
+    reveal_type(x)  # revealed: int & ~Literal[1]
+elif x != 3:
+    reveal_type(x)  # revealed: int & ~Literal[1] & ~Literal[2] & ~Literal[3]
+```
+
+## Positive contributions become negative in elif-else blocks, with simplification
+
+```py
+def bool_instance() -> bool:
+    return True
+
+x = 1 if bool_instance() else 2 if bool_instance() else 3
+
+if x == 1:
+    # TODO should be Literal[1]
+    reveal_type(x)  # revealed: Literal[1, 2, 3]
+elif x == 2:
+    # TODO should be Literal[2]
+    reveal_type(x)  # revealed: Literal[2, 3]
+else:
+    reveal_type(x)  # revealed: Literal[3]
+```
+
+## Multiple negative contributions using elif, with simplification
+
+```py
+def bool_instance() -> bool:
+    return True
+
+x = 1 if bool_instance() else 2 if bool_instance() else 3
+
+if x != 1:
+    reveal_type(x)  # revealed: Literal[2, 3]
+elif x != 2:
+    # TODO should be `Literal[1]`
+    reveal_type(x)  # revealed: Literal[1, 3]
+elif x == 3:
+    # TODO should be Never
+    reveal_type(x)  # revealed: Literal[1, 2, 3]
+else:
+    # TODO should be Never
+    reveal_type(x)  # revealed: Literal[1, 2]
+```

crates/red_knot_python_semantic/resources/mdtest/narrow/conditionals_is.md

@@ -11,6 +11,8 @@ x = None if flag else 1
 
 if x is None:
     reveal_type(x)  # revealed: None
+else:
+    reveal_type(x)  # revealed: Literal[1]
 
 reveal_type(x)  # revealed: None | Literal[1]
 ```
@@ -30,6 +32,8 @@ y = x if flag else None
 
 if y is x:
     reveal_type(y)  # revealed: A
+else:
+    reveal_type(y)  # revealed: A | None
 
 reveal_type(y)  # revealed: A | None
 ```
@@ -50,4 +54,26 @@ reveal_type(y)  # revealed: bool
 if y is x is False:  # Interpreted as `(y is x) and (x is False)`
     reveal_type(x)  # revealed: Literal[False]
     reveal_type(y)  # revealed: bool
+else:
+    # The negation of the clause above is (y is not x) or (x is not False)
+    # So we can't narrow the type of x or y here, because each arm of the `or` could be true
+    reveal_type(x)  # revealed: bool
+    reveal_type(y)  # revealed: bool
+```
+
+## `is` in elif clause
+
+```py
+def bool_instance() -> bool:
+    return True
+
+x = None if bool_instance() else (1 if bool_instance() else True)
+
+reveal_type(x)  # revealed: None | Literal[1] | Literal[True]
+if x is None:
+    reveal_type(x)  # revealed: None
+elif x is True:
+    reveal_type(x)  # revealed: Literal[True]
+else:
+    reveal_type(x)  # revealed: Literal[1]
 ```

crates/red_knot_python_semantic/resources/mdtest/narrow/conditionals_is_not.md

@@ -13,6 +13,8 @@ x = None if flag else 1
 
 if x is not None:
     reveal_type(x)  # revealed: Literal[1]
+else:
+    reveal_type(x)  # revealed: None
 
 reveal_type(x)  # revealed: None | Literal[1]
 ```
@@ -29,6 +31,8 @@ reveal_type(x)  # revealed: bool
 
 if x is not False:
     reveal_type(x)  # revealed: Literal[True]
+else:
+    reveal_type(x)  # revealed: Literal[False]
 ```
 
 ## `is not` for non-singleton types
@@ -43,6 +47,27 @@ y = 345
 
 if x is not y:
     reveal_type(x)  # revealed: Literal[345]
+else:
+    reveal_type(x)  # revealed: Literal[345]
+```
+
+## `is not` for other types
+
+```py
+def bool_instance() -> bool:
+    return True
+
+class A: ...
+
+x = A()
+y = x if bool_instance() else None
+
+if y is not x:
+    reveal_type(y)  # revealed: A | None
+else:
+    reveal_type(y)  # revealed: A
+
+reveal_type(y)  # revealed: A | None
 ```
 
 ## `is not` in chained comparisons
@@ -63,4 +88,10 @@ reveal_type(y)  # revealed: bool
 if y is not x is not False:  # Interpreted as `(y is not x) and (x is not False)`
     reveal_type(x)  # revealed: Literal[True]
     reveal_type(y)  # revealed: bool
+else:
+    # The negation of the clause above is (y is x) or (x is False)
+    # So we can't narrow the type of x or y here, because each arm of the `or` could be true
+
+    reveal_type(x)  # revealed: bool
+    reveal_type(y)  # revealed: bool
 ```

crates/red_knot_python_semantic/resources/mdtest/narrow/conditionals_nested.md

@@ -3,7 +3,8 @@
 ## Multiple negative contributions
 
 ```py
-def int_instance() -> int: ...
+def int_instance() -> int:
+    return 42
 
 x = int_instance()
 
@@ -27,3 +28,29 @@ if x != 1:
     if x != 2:
         reveal_type(x)  # revealed: Literal[3]
 ```
+
+## elif-else blocks
+
+```py
+def bool_instance() -> bool:
+    return True
+
+x = 1 if bool_instance() else 2 if bool_instance() else 3
+
+if x != 1:
+    reveal_type(x)  # revealed: Literal[2, 3]
+    if x == 2:
+        # TODO should be `Literal[2]`
+        reveal_type(x)  # revealed: Literal[2, 3]
+    elif x == 3:
+        reveal_type(x)  # revealed: Literal[3]
+    else:
+        reveal_type(x)  # revealed: Never
+
+elif x != 2:
+    # TODO should be Literal[1]
+    reveal_type(x)  # revealed: Literal[1, 3]
+else:
+    # TODO should be Never
+    reveal_type(x)  # revealed: Literal[1, 2, 3]
+```

crates/red_knot_python_semantic/resources/mdtest/narrow/conditionals_not_eq.md

@@ -11,6 +11,9 @@ x = None if flag else 1
 
 if x != None:
     reveal_type(x)  # revealed: Literal[1]
+else:
+    # TODO should be None
+    reveal_type(x)  # revealed: None | Literal[1]
 ```
 
 ## `!=` for other singleton types
@@ -24,6 +27,9 @@ x = True if flag else False
 
 if x != False:
     reveal_type(x)  # revealed: Literal[True]
+else:
+    # TODO should be Literal[False]
+    reveal_type(x)  # revealed: bool
 ```
 
 ## `x != y` where `y` is of literal type
@@ -54,6 +60,25 @@ C = A if flag else B
 
 if C != A:
     reveal_type(C)  # revealed: Literal[B]
+else:
+    # TODO should be Literal[A]
+    reveal_type(C)  # revealed: Literal[A, B]
+```
+
+## `x != y` where `y` has multiple single-valued options
+
+```py
+def bool_instance() -> bool:
+    return True
+
+x = 1 if bool_instance() else 2
+y = 2 if bool_instance() else 3
+
+if x != y:
+    reveal_type(x)  # revealed: Literal[1, 2]
+else:
+    # TODO should be Literal[2]
+    reveal_type(x)  # revealed: Literal[1, 2]
 ```
 
 ## `!=` for non-single-valued types
@@ -74,3 +99,21 @@ y = int_instance()
 if x != y:
     reveal_type(x)  # revealed: int | None
 ```
+
+## Mix of single-valued and non-single-valued types
+
+```py
+def int_instance() -> int:
+    return 42
+
+def bool_instance() -> bool:
+    return True
+
+x = 1 if bool_instance() else 2
+y = 2 if bool_instance() else int_instance()
+
+if x != y:
+    reveal_type(x)  # revealed: Literal[1, 2]
+else:
+    reveal_type(x)  # revealed: Literal[1, 2]
+```

crates/red_knot_python_semantic/resources/mdtest/narrow/isinstance.md

@@ -40,6 +40,8 @@ x = 1 if flag else "a"
 
 if isinstance(x, (int, str)):
     reveal_type(x)  # revealed: Literal[1] | Literal["a"]
+else:
+    reveal_type(x)  # revealed: Never
 
 if isinstance(x, (int, bytes)):
     reveal_type(x)  # revealed: Literal[1]
@@ -51,6 +53,8 @@ if isinstance(x, (bytes, str)):
 # one of the possibilities:
 if isinstance(x, (int, object)):
     reveal_type(x)  # revealed: Literal[1] | Literal["a"]
+else:
+    reveal_type(x)  # revealed: Never
 
 y = 1 if flag1 else "a" if flag2 else b"b"
 if isinstance(y, (int, str)):
@@ -75,13 +79,16 @@ x = 1 if flag else "a"
 
 if isinstance(x, (bool, (bytes, int))):
     reveal_type(x)  # revealed: Literal[1]
+else:
+    reveal_type(x)  # revealed: Literal["a"]
 ```
 
 ## Class types
 
 ```py
 class A: ...
 class B: ...
+class C: ...
 
 def get_object() -> object: ...
 
@@ -91,6 +98,16 @@ if isinstance(x, A):
     reveal_type(x)  # revealed: A
     if isinstance(x, B):
         reveal_type(x)  # revealed: A & B
+    else:
+        reveal_type(x)  # revealed: A & ~B
+
+if isinstance(x, (A, B)):
+    reveal_type(x)  # revealed: A | B
+elif isinstance(x, (A, C)):
+    reveal_type(x)  # revealed: C & ~A & ~B
+else:
+    # TODO: Should be simplified to ~A & ~B & ~C
+    reveal_type(x)  # revealed: object & ~A & ~B & ~C
 ```
 
 ## No narrowing for instances of `builtins.type`

crates/red_knot_python_semantic/resources/mdtest/subscript/bytes.md

@@ -26,7 +26,8 @@ reveal_type(y)  # revealed: Unknown
 ## Function return
 
 ```py
-def int_instance() -> int: ...
+def int_instance() -> int:
+    return 42
 
 a = b"abcde"[int_instance()]
 # TODO: Support overloads... Should be `bytes`

crates/red_knot_python_semantic/resources/mdtest/subscript/string.md

@@ -23,7 +23,8 @@ reveal_type(b)  # revealed: Unknown
 ## Function return
 
 ```py
-def int_instance() -> int: ...
+def int_instance() -> int:
+    return 42
 
 a = "abcde"[int_instance()]
 # TODO: Support overloads... Should be `str`