Is a potential update to the type-inference safe involving functions/lambda that return a variable from an outer scope?

[Daraan]

I inherited some code that repeatedly uses different class wrappers, but in the end it is the same type no matter which branch was taken, it looks for example like this:

def create():
    a = Var(1)  # a Var
    if foo:
        a = Foo(a)  # now a Foo
    else:
        a = Baz(a)  # now a Baz
    a = Unified(a)  # now Unified
    reveal_type(a)  # Unified
    A = [lambda: a for _ in range(3)]
    reveal_type(A)  # list[() -> Var | Foo | Baz | Unified]
    return A

The inferred type for A and the return type is list[() -> Var | Foo | Baz | Unified].
This comes a bit as a surprise as a does not change and is always Unified afterwards.
Of course the value of a could change after the definition of A and the lambdas would return the new value.
Nevertheless, I think the type of A could be improved to be:
reveal_type(A) # list[() -> Unified | <aTypesAfterDefinitionOfA>*]
instead of being:
reveal_type(A) # list[() -> <aTypesBeforeDefinitionOfA>* | Unified | <aTypesAfterDefinitionOfA>*]
As long as a is not global/nonlocal.

I assume this needs some backward/forward checking but in the end would avoid false types being inferred.
I wonder if such a change is feasible and could be an enhancement request or if something could go very wrong.

---

Below I included a toy example:

Code sample in pyright playground

from typing import reveal_type

def foo():
    a = 1
    a = "hello" + str(a)
    a = (1, 2, a)
    reveal_type(a)  # tuple[Literal[1], Literal[2], str]
    
    # Use a in function/lambda
    A = [lambda: a for _ in range(3)]
    reveal_type(A)  # "list[() -> (int | str | tuple[Literal[1], Literal[2], str] | set[int])]".
    
    def bar():
        return a
    reveal_type(bar)# "() -> (int | str | tuple[Literal[1], Literal[2], str] | set[int])".
    
    # Change type of a that would affect A / bar
    a = {1}
    reveal_type(a)  # set[int]
    return A

reveal_type(foo())  # "list[() -> (int | str | tuple[Literal[1], Literal[2], str] | set[int])]".

[erictraut]

In general, locally-narrowed types cannot be used for variables captured by an inner scope. Pyright is able to use the locally-narrowed type under very specific circumstances where it can prove that the value of that variable does not change after the inner-scope definition (a lambda expression or def statement). This check is very complex and expensive. For more details, refer to this documentation.

In your toy sample above, you're violating the requirements that make it safe to use the locally-narrowed type, so pyright cannot use it in that case. It would normally fall back on the declared type of the variable in this case, but you haven't declared the type of a in your example, so it evaluates the union of all assigned types.

I think this is a case where you're going to need to modify your code if you want it to work better with static type checking. Perhaps add some type declarations or change the names of some variables to avoid reusing the same variable name for different uses within the same scope.