[ty] basic narrowing on attribute and subscript expressions

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

@@ -0,0 +1,209 @@
+# Narrowing for complex targets (attribute expressions, subscripts)
+
+We support type narrowing for attributes and subscripts.
+
+## Attribute narrowing
+
+### Basic
+
+```py
+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
+```
+
+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

@@ -134,8 +134,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: D | None
 
 a.b.c1 = C()
 a.b.c1.d = D()
@@ -173,12 +172,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 +188,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 +236,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 +290,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

@@ -18,4 +18,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

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

crates/ty_python_semantic/src/place.rs

@@ -661,6 +661,7 @@ fn place_by_id<'db>(
             // See mdtest/known_constants.md#user-defined-type_checking for details.
             let is_considered_non_modifiable = place_table(db, scope)
                 .place_expr(place_id)
+                .expr
                 .is_name_and(|name| matches!(name, "__slots__" | "TYPE_CHECKING"));
 
             if scope.file(db).is_stub(db.upcast()) {
@@ -1124,8 +1125,8 @@ mod implicit_globals {
 
         module_type_symbol_table
             .places()
-            .filter(|symbol| symbol.is_declared() && symbol.is_name())
-            .map(semantic_index::place::PlaceExpr::expect_name)
+            .filter(|place| place.is_declared() && place.is_name())
+            .map(semantic_index::place::PlaceExprWithFlags::expect_name)
             .filter(|symbol_name| {
                 !matches!(&***symbol_name, "__dict__" | "__getattr__" | "__init__")
             })

crates/ty_python_semantic/src/semantic_index.rs

@@ -37,8 +37,8 @@ mod use_def;
 mod visibility_constraints;
 
 pub(crate) use self::use_def::{
-    BindingWithConstraints, BindingWithConstraintsIterator, DeclarationWithConstraint,
-    DeclarationsIterator,
+    ApplicableConstraints, BindingWithConstraints, BindingWithConstraintsIterator,
+    DeclarationWithConstraint, DeclarationsIterator,
 };
 
 type PlaceSet = hashbrown::HashMap<ScopedPlaceId, (), FxBuildHasher>;