[ty] basic narrowing on attribute and subscript expressions (#17643)

## Summary

This PR closes astral-sh/ty#164.

This PR introduces a basic type narrowing mechanism for
attribute/subscript expressions.
Member accesses, int literal subscripts, string literal subscripts are
supported (same as mypy and pyright).

## Test Plan

New test cases are added to `mdtest/narrow/complex_target.md`.

---------

Co-authored-by: David Peter <mail@david-peter.de>

crates/ty_python_semantic/resources/mdtest/attributes.md

@@ -751,7 +751,8 @@ reveal_type(C.pure_class_variable)  # revealed: Unknown
 # and the assignment is properly attributed to the class method.
 # error: [invalid-attribute-access] "Cannot assign to instance attribute `pure_class_variable` from the class object `<class 'C'>`"
 C.pure_class_variable = "overwritten on class"
-
+# TODO: should be no error
+# error: [unresolved-attribute] "Attribute `pure_class_variable` can only be accessed on instances, not on the class object `<class 'C'>` itself."
 reveal_type(C.pure_class_variable)  # revealed: Literal["overwritten on class"]
 
 c_instance = C()

crates/ty_python_semantic/resources/mdtest/narrow/assignment.md

@@ -87,6 +87,12 @@ class _:
         reveal_type(a.y)  # revealed: Unknown | None
         reveal_type(a.z)  # revealed: Unknown | None
 
+a = A()
+# error: [unresolved-attribute]
+a.dynamically_added = 0
+# error: [unresolved-attribute]
+reveal_type(a.dynamically_added)  # revealed: Literal[0]
+
 # error: [unresolved-reference]
 does.nt.exist = 0
 # error: [unresolved-reference]

crates/ty_python_semantic/resources/mdtest/narrow/complex_target.md

@@ -0,0 +1,224 @@
+# Narrowing for complex targets (attribute expressions, subscripts)
+
+We support type narrowing for attributes and subscripts.
+
+## Attribute narrowing
+
+### Basic
+
+```py
+from ty_extensions import Unknown
+
+class C:
+    x: int | None = None
+
+c = C()
+
+reveal_type(c.x)  # revealed: int | None
+
+if c.x is not None:
+    reveal_type(c.x)  # revealed: int
+else:
+    reveal_type(c.x)  # revealed: None
+
+if c.x is not None:
+    c.x = None
+
+reveal_type(c.x)  # revealed: None
+
+c = C()
+
+if c.x is None:
+    c.x = 1
+
+reveal_type(c.x)  # revealed: int
+
+class _:
+    reveal_type(c.x)  # revealed: int
+
+c = C()
+
+class _:
+    if c.x is None:
+        c.x = 1
+    reveal_type(c.x)  # revealed: int
+
+# TODO: should be `int`
+reveal_type(c.x)  # revealed: int | None
+
+class D:
+    x = None
+
+def unknown() -> Unknown:
+    return 1
+
+d = D()
+reveal_type(d.x)  # revealed: Unknown | None
+d.x = 1
+reveal_type(d.x)  # revealed: Literal[1]
+d.x = unknown()
+reveal_type(d.x)  # revealed: Unknown
+```
+
+Narrowing can be "reset" by assigning to the attribute:
+
+```py
+c = C()
+
+if c.x is None:
+    reveal_type(c.x)  # revealed: None
+    c.x = 1
+    reveal_type(c.x)  # revealed: Literal[1]
+    c.x = None
+    reveal_type(c.x)  # revealed: None
+
+reveal_type(c.x)  # revealed: int | None
+```
+
+Narrowing can also be "reset" by assigning to the object:
+
+```py
+c = C()
+
+if c.x is None:
+    reveal_type(c.x)  # revealed: None
+    c = C()
+    reveal_type(c.x)  # revealed: int | None
+
+reveal_type(c.x)  # revealed: int | None
+```
+
+### Multiple predicates
+
+```py
+class C:
+    value: str | None
+
+def foo(c: C):
+    if c.value and len(c.value):
+        reveal_type(c.value)  # revealed: str & ~AlwaysFalsy
+
+    # error: [invalid-argument-type] "Argument to function `len` is incorrect: Expected `Sized`, found `str | None`"
+    if len(c.value) and c.value:
+        reveal_type(c.value)  # revealed: str & ~AlwaysFalsy
+
+    if c.value is None or not len(c.value):
+        reveal_type(c.value)  # revealed: str | None
+    else:  # c.value is not None and len(c.value)
+        # TODO: should be # `str & ~AlwaysFalsy`
+        reveal_type(c.value)  # revealed: str
+```
+
+### Generic class
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+```py
+class C[T]:
+    x: T
+    y: T
+
+    def __init__(self, x: T):
+        self.x = x
+        self.y = x
+
+def f(a: int | None):
+    c = C(a)
+    reveal_type(c.x)  # revealed: int | None
+    reveal_type(c.y)  # revealed: int | None
+    if c.x is not None:
+        reveal_type(c.x)  # revealed: int
+        # In this case, it may seem like we can narrow it down to `int`,
+        # but different values ​​may be reassigned to `x` and `y` in another place.
+        reveal_type(c.y)  # revealed: int | None
+
+def g[T](c: C[T]):
+    reveal_type(c.x)  # revealed: T
+    reveal_type(c.y)  # revealed: T
+    reveal_type(c)  # revealed: C[T]
+
+    if isinstance(c.x, int):
+        reveal_type(c.x)  # revealed: T & int
+        reveal_type(c.y)  # revealed: T
+        reveal_type(c)  # revealed: C[T]
+    if isinstance(c.x, int) and isinstance(c.y, int):
+        reveal_type(c.x)  # revealed: T & int
+        reveal_type(c.y)  # revealed: T & int
+        # TODO: Probably better if inferred as `C[T & int]` (mypy and pyright don't support this)
+        reveal_type(c)  # revealed: C[T]
+```
+
+### With intermediate scopes
+
+```py
+class C:
+    def __init__(self):
+        self.x: int | None = None
+        self.y: int | None = None
+
+c = C()
+reveal_type(c.x)  # revealed: int | None
+if c.x is not None:
+    reveal_type(c.x)  # revealed: int
+    reveal_type(c.y)  # revealed: int | None
+
+if c.x is not None:
+    def _():
+        reveal_type(c.x)  # revealed: Unknown | int | None
+
+def _():
+    if c.x is not None:
+        reveal_type(c.x)  # revealed: (Unknown & ~None) | int
+```
+
+## Subscript narrowing
+
+### Number subscript
+
+```py
+def _(t1: tuple[int | None, int | None], t2: tuple[int, int] | tuple[None, None]):
+    if t1[0] is not None:
+        reveal_type(t1[0])  # revealed: int
+        reveal_type(t1[1])  # revealed: int | None
+
+    n = 0
+    if t1[n] is not None:
+        # Non-literal subscript narrowing are currently not supported, as well as mypy, pyright
+        reveal_type(t1[0])  # revealed: int | None
+        reveal_type(t1[n])  # revealed: int | None
+        reveal_type(t1[1])  # revealed: int | None
+
+    if t2[0] is not None:
+        # TODO: should be int
+        reveal_type(t2[0])  # revealed: Unknown & ~None
+        # TODO: should be int
+        reveal_type(t2[1])  # revealed: Unknown
+```
+
+### String subscript
+
+```py
+def _(d: dict[str, str | None]):
+    if d["a"] is not None:
+        reveal_type(d["a"])  # revealed: str
+        reveal_type(d["b"])  # revealed: str | None
+```
+
+## Combined attribute and subscript narrowing
+
+```py
+class C:
+    def __init__(self):
+        self.x: tuple[int | None, int | None] = (None, None)
+
+class D:
+    def __init__(self):
+        self.c: tuple[C] | None = None
+
+d = D()
+if d.c is not None and d.c[0].x[0] is not None:
+    reveal_type(d.c[0].x[0])  # revealed: int
+```

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md

@@ -135,7 +135,7 @@ class _:
         class _3:
             reveal_type(a)  # revealed: A
             # TODO: should be `D | None`
-            reveal_type(a.b.c1.d)  # revealed: D
+            reveal_type(a.b.c1.d)  # revealed: Unknown
 
 a.b.c1 = C()
 a.b.c1.d = D()
@@ -173,12 +173,10 @@ def f(x: str | None):
             reveal_type(g)  # revealed: str
 
         if a.x is not None:
-            # TODO(#17643): should be `Unknown | str`
-            reveal_type(a.x)  # revealed: Unknown | str | None
+            reveal_type(a.x)  # revealed: (Unknown & ~None) | str
 
         if l[0] is not None:
-            # TODO(#17643): should be `str`
-            reveal_type(l[0])  # revealed: str | None
+            reveal_type(l[0])  # revealed: str
 
     class C:
         if x is not None:
@@ -191,12 +189,10 @@ def f(x: str | None):
             reveal_type(g)  # revealed: str
 
         if a.x is not None:
-            # TODO(#17643): should be `Unknown | str`
-            reveal_type(a.x)  # revealed: Unknown | str | None
+            reveal_type(a.x)  # revealed: (Unknown & ~None) | str
 
         if l[0] is not None:
-            # TODO(#17643): should be `str`
-            reveal_type(l[0])  # revealed: str | None
+            reveal_type(l[0])  # revealed: str
 
     # TODO: should be str
     # This could be fixed if we supported narrowing with if clauses in comprehensions.
@@ -241,22 +237,18 @@ def f(x: str | None):
             reveal_type(a.x)  # revealed: Unknown | str | None
 
         class D:
-            # TODO(#17643): should be `Unknown | str`
-            reveal_type(a.x)  # revealed: Unknown | str | None
+            reveal_type(a.x)  # revealed: (Unknown & ~None) | str
 
-        # TODO(#17643): should be `Unknown | str`
-        [reveal_type(a.x) for _ in range(1)]  # revealed: Unknown | str | None
+        [reveal_type(a.x) for _ in range(1)]  # revealed: (Unknown & ~None) | str
 
     if l[0] is not None:
         def _():
             reveal_type(l[0])  # revealed: str | None
 
         class D:
-            # TODO(#17643): should be `str`
-            reveal_type(l[0])  # revealed: str | None
+            reveal_type(l[0])  # revealed: str
 
-        # TODO(#17643): should be `str`
-        [reveal_type(l[0]) for _ in range(1)]  # revealed: str | None
+        [reveal_type(l[0]) for _ in range(1)]  # revealed: str
 ```
 
 ### Narrowing constraints introduced in multiple scopes
@@ -299,24 +291,20 @@ def f(x: str | Literal[1] | None):
         if a.x is not None:
             def _():
                 if a.x != 1:
-                    # TODO(#17643): should be `Unknown | str | None`
-                    reveal_type(a.x)  # revealed: Unknown | str | Literal[1] | None
+                    reveal_type(a.x)  # revealed: (Unknown & ~Literal[1]) | str | None
 
             class D:
                 if a.x != 1:
-                    # TODO(#17643): should be `Unknown | str`
-                    reveal_type(a.x)  # revealed: Unknown | str | Literal[1] | None
+                    reveal_type(a.x)  # revealed: (Unknown & ~Literal[1] & ~None) | str
 
         if l[0] is not None:
             def _():
                 if l[0] != 1:
-                    # TODO(#17643): should be `str | None`
-                    reveal_type(l[0])  # revealed: str | Literal[1] | None
+                    reveal_type(l[0])  # revealed: str | None
 
             class D:
                 if l[0] != 1:
-                    # TODO(#17643): should be `str`
-                    reveal_type(l[0])  # revealed: str | Literal[1] | None
+                    reveal_type(l[0])  # revealed: str
 ```
 
 ### Narrowing constraints with bindings in class scope, and nested scopes

crates/ty_python_semantic/resources/mdtest/narrow/type_guards.md

@@ -220,8 +220,7 @@ def _(a: tuple[str, int] | tuple[int, str], c: C[Any]):
     if reveal_type(is_int(a[0])):  # revealed: TypeIs[int @ a[0]]
         # TODO: Should be `tuple[int, str]`
         reveal_type(a)  # revealed: tuple[str, int] | tuple[int, str]
-        # TODO: Should be `int`
-        reveal_type(a[0])  # revealed: Unknown
+        reveal_type(a[0])  # revealed: Unknown & int
 
     # TODO: Should be `TypeGuard[str @ c.v]`
     if reveal_type(guard_str(c.v)):  # revealed: @Todo(`TypeGuard[]` special form)
@@ -231,8 +230,7 @@ def _(a: tuple[str, int] | tuple[int, str], c: C[Any]):
 
     if reveal_type(is_int(c.v)):  # revealed: TypeIs[int @ c.v]
         reveal_type(c)  # revealed: C[Any]
-        # TODO: Should be `int`
-        reveal_type(c.v)  # revealed: Any
+        reveal_type(c.v)  # revealed: Any & int
 ```
 
 Indirect usage is supported within the same scope:

crates/ty_python_semantic/resources/primer/bad.txt

@@ -17,4 +17,5 @@ setuptools  # vendors packaging, see above
 spack  # slow, success, but mypy-primer hangs processing the output
 spark  # too many iterations
 steam.py  # hangs (single threaded)
+tornado # bad use-def map (https://github.com/astral-sh/ty/issues/365)
 xarray  # too many iterations

crates/ty_python_semantic/resources/primer/good.txt

@@ -110,7 +110,6 @@ stone
 strawberry
 streamlit
 sympy
-tornado
 trio
 twine
 typeshed-stats