[ty] improve lazy scope place lookup (#19321)

Co-authored-by: David Peter <sharkdp@users.noreply.github.com>
Co-authored-by: Carl Meyer <carl@oddbird.net>
Co-authored-by: Micha Reiser <micha@reiser.io>

Author: 4 people

crates/ty_python_semantic/resources/mdtest/del.md

@@ -123,7 +123,7 @@ def enclosing():
         nonlocal x
         def bar():
             # allowed, refers to `x` in `enclosing`
-            reveal_type(x)  # revealed: Unknown | Literal[2]
+            reveal_type(x)  # revealed: Literal[2]
         bar()
         del x  # allowed, deletes `x` in `enclosing` (though we don't track that)
 ```

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

@@ -238,6 +238,69 @@ def f(x: str | None):
 
         [reveal_type(x) for _ in range(1)]  # revealed: str
 
+    # When there is a reassignment, any narrowing constraints on the place are invalidated in lazy scopes.
+    x = None
+```
+
+If a variable defined in a private scope is never reassigned, narrowing remains in effect in the
+inner lazy scope.
+
+```py
+def f(const: str | None):
+    if const is not None:
+        def _():
+            # The `const is not None` narrowing constraint is still valid since `const` has not been reassigned
+            reveal_type(const)  # revealed: str
+
+        class C2:
+            reveal_type(const)  # revealed: str
+
+        [reveal_type(const) for _ in range(1)]  # revealed: str
+```
+
+And even if there is an attribute or subscript assignment to the variable, narrowing of the variable
+is still valid in the inner lazy scope.
+
+```py
+def f(l: list[str | None] | None):
+    if l is not None:
+        def _():
+            reveal_type(l)  # revealed: list[str | None]
+        l[0] = None
+
+def f(a: A):
+    if a:
+        def _():
+            reveal_type(a)  # revealed: A & ~AlwaysFalsy
+    a.x = None
+```
+
+Narrowing is invalidated if a `nonlocal` declaration is made within a lazy scope.
+
+```py
+def f(non_local: str | None):
+    if non_local is not None:
+        def _():
+            nonlocal non_local
+            non_local = None
+
+        def _():
+            reveal_type(non_local)  # revealed: str | None
+
+def f(non_local: str | None):
+    def _():
+        nonlocal non_local
+        non_local = None
+    if non_local is not None:
+        def _():
+            reveal_type(non_local)  # revealed: str | None
+```
+
+The same goes for public variables, attributes, and subscripts, because it is difficult to track all
+of their changes.
+
+```py
+def f():
     if g is not None:
         def _():
             reveal_type(g)  # revealed: str | None
@@ -249,6 +312,7 @@ def f(x: str | None):
 
     if a.x is not None:
         def _():
+            # Lazy nested scope narrowing is not performed on attributes/subscripts because it's difficult to track their changes.
             reveal_type(a.x)  # revealed: Unknown | str | None
 
         class D:
@@ -282,7 +346,7 @@ l: list[str | Literal[1] | None] = [None]
 
 def f(x: str | Literal[1] | None):
     class C:
-        if x is not None:
+        if x is not None:  # TODO: should be an unresolved-reference error
             def _():
                 if x != 1:
                     reveal_type(x)  # revealed: str | None
@@ -293,6 +357,38 @@ def f(x: str | Literal[1] | None):
 
             [reveal_type(x) for _ in range(1) if x != 1]  # revealed: str
 
+        x = None
+
+    def _():
+        # error: [unresolved-reference]
+        if x is not None:
+            def _():
+                if x != 1:
+                    reveal_type(x)  # revealed: Never
+        x = None
+
+def f(const: str | Literal[1] | None):
+    class C:
+        if const is not None:
+            def _():
+                if const != 1:
+                    # TODO: should be `str`
+                    reveal_type(const)  # revealed: str | None
+
+            class D:
+                if const != 1:
+                    reveal_type(const)  # revealed: str
+
+            [reveal_type(const) for _ in range(1) if const != 1]  # revealed: str
+
+    def _():
+        if const is not None:
+            def _():
+                if const != 1:
+                    reveal_type(const)  # revealed: str
+
+def f():
+    class C:
         if g is not None:
             def _():
                 if g != 1:

crates/ty_python_semantic/resources/mdtest/public_types.md

@@ -20,10 +20,7 @@ def outer() -> None:
     x = A()
 
     def inner() -> None:
-        # TODO: We might ideally be able to eliminate `Unknown` from the union here since `x` resolves to an
-        # outer scope that is a function scope (as opposed to module global scope), and `x` is never declared
-        # nonlocal in a nested scope that also assigns to it.
-        reveal_type(x)  # revealed: Unknown | A | B
+        reveal_type(x)  # revealed: A | B
     # This call would observe `x` as `A`.
     inner()
 
@@ -40,7 +37,7 @@ def outer(flag: bool) -> None:
     x = A()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: Unknown | A | B | C
+        reveal_type(x)  # revealed: A | B | C
     inner()
 
     if flag:
@@ -62,7 +59,7 @@ def outer() -> None:
     x = A()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: Unknown | A | C
+        reveal_type(x)  # revealed: A | C
     inner()
 
     if False:
@@ -76,7 +73,7 @@ def outer(flag: bool) -> None:
     x = A()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: Unknown | A | C
+        reveal_type(x)  # revealed: A | C
     inner()
 
     if flag:
@@ -96,24 +93,18 @@ def outer(flag: bool) -> None:
         x = A()
 
         def inner() -> None:
-            reveal_type(x)  # revealed: Unknown | A
+            reveal_type(x)  # revealed: A
         inner()
 ```
 
-In the future, we may try to be smarter about which bindings must or must not be a visible to a
-given nested scope, depending where it is defined. In the above case, this shouldn't change the
-behavior -- `x` is defined before `inner` in the same branch, so should be considered
-definitely-bound for `inner`. But in other cases we may want to emit `possibly-unresolved-reference`
-in future:
-
 ```py
 def outer(flag: bool) -> None:
     if flag:
         x = A()
 
     def inner() -> None:
         # TODO: Ideally, we would emit a possibly-unresolved-reference error here.
-        reveal_type(x)  # revealed: Unknown | A
+        reveal_type(x)  # revealed: A
     inner()
 ```
 
@@ -126,7 +117,7 @@ def outer() -> None:
     x = A()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: Unknown | A
+        reveal_type(x)  # revealed: A
     inner()
 
     return
@@ -136,7 +127,7 @@ def outer(flag: bool) -> None:
     x = A()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: Unknown | A | B
+        reveal_type(x)  # revealed: A | B
     if flag:
         x = B()
         inner()
@@ -161,7 +152,7 @@ def f0() -> None:
         def f2() -> None:
             def f3() -> None:
                 def f4() -> None:
-                    reveal_type(x)  # revealed: Unknown | A | B
+                    reveal_type(x)  # revealed: A | B
                 f4()
             f3()
         f2()
@@ -172,6 +163,29 @@ def f0() -> None:
     f1()
 ```
 
+## Narrowing
+
+In general, it is not safe to narrow the public type of a symbol using constraints introduced in an
+outer scope (because the symbol's value may have changed by the time the lazy scope is actually
+evaluated), but they can be applied if there is no reassignment of the symbol.
+
+```py
+class A: ...
+
+def outer(x: A | None):
+    if x is not None:
+        def inner() -> None:
+            reveal_type(x)  # revealed: A | None
+        inner()
+    x = None
+
+def outer(x: A | None):
+    if x is not None:
+        def inner() -> None:
+            reveal_type(x)  # revealed: A
+        inner()
+```
+
 ## At module level
 
 The behavior is the same if the outer scope is the global scope of a module:
@@ -232,32 +246,16 @@ def _():
 
 ## Limitations
 
-### Type narrowing
-
-We currently do not further analyze control flow, so we do not support cases where the inner scope
-is only executed in a branch where the type of `x` is narrowed:
-
-```py
-class A: ...
-
-def outer(x: A | None):
-    if x is not None:
-        def inner() -> None:
-            # TODO: should ideally be `A`
-            reveal_type(x)  # revealed: A | None
-        inner()
-```
-
 ### Shadowing
 
-Similarly, since we do not analyze control flow in the outer scope here, we assume that `inner()`
-could be called between the two assignments to `x`:
+Since we do not analyze control flow in the outer scope here, we assume that `inner()` could be
+called between the two assignments to `x`:
 
 ```py
 def outer() -> None:
     def inner() -> None:
-        # TODO: this should ideally be `Unknown | Literal[1]`, but no other type checker supports this either
-        reveal_type(x)  # revealed: Unknown | None | Literal[1]
+        # TODO: this should ideally be `Literal[1]`, but no other type checker supports this either
+        reveal_type(x)  # revealed: None | Literal[1]
     x = None
 
     # [additional code here]
@@ -279,8 +277,8 @@ def outer() -> None:
     x = 1
 
     def inner() -> None:
-        # TODO: this should be `Unknown | Literal[1]`. Mypy and pyright support this.
-        reveal_type(x)  # revealed: Unknown | None | Literal[1]
+        # TODO: this should be `Literal[1]`. Mypy and pyright support this.
+        reveal_type(x)  # revealed: None | Literal[1]
     inner()
 ```
 
@@ -314,8 +312,8 @@ def outer() -> None:
     set_x()
 
     def inner() -> None:
-        # TODO: this should ideally be `Unknown | None | Literal[1]`. Mypy and pyright support this.
-        reveal_type(x)  # revealed: Unknown | None
+        # TODO: this should ideally be `None | Literal[1]`. Mypy and pyright support this.
+        reveal_type(x)  # revealed: None
     inner()
 ```
 

crates/ty_python_semantic/resources/mdtest/scopes/eager.md

@@ -299,7 +299,7 @@ def _():
     x = 1
 
     def f():
-        # revealed: Unknown | Literal[1, 2]
+        # revealed: Literal[1, 2]
         [reveal_type(x) for a in range(1)]
     x = 2
 ```
@@ -316,7 +316,7 @@ def _():
 
     class A:
         def f():
-            # revealed: Unknown | Literal[1, 2]
+            # revealed: Literal[1, 2]
             reveal_type(x)
 
     x = 2
@@ -333,7 +333,7 @@ def _():
 
     def f():
         def g():
-            # revealed: Unknown | Literal[1, 2]
+            # revealed: Literal[1, 2]
             reveal_type(x)
     x = 2
 ```
@@ -351,7 +351,7 @@ def _():
 
     class A:
         def f():
-            # revealed: Unknown | Literal[1, 2]
+            # revealed: Literal[1, 2]
             [reveal_type(x) for a in range(1)]
 
     x = 2