Add flag to allow more flexible variable redefinition (#18727)

Infer union types for simple variables from multiple assignments, if the
variable isn't annotated. The feature is enabled via
`--allow-redefinition-new`. `--local-partial-types` must also be
enabled.

This is still experimental and has known issues, so it's not documented
anywhere. It works well enough that it can be used for non-trivial
experimentation, however.

Closes #6233. Closes #6232. Closes #18568. Fixes #18619.

In this example, the type of `x` is inferred as `int | str` when using
the new behavior:
```py
def f(i: int, s : str) -> int | str:
    if i > 5:
        x = i
    else:
        x = s  # No longer an error
    reveal_type(x)  # int | str
    return s
```

Here is a summary of how it works:
* Assignment widens the inferred type of a variable and always narrows
(when there is no annotation).
* Simple variable lvalues are put into the binder on initial assignment
when using the new feature. We need to be able to track whether a
variable is defined or not to infer correct types (see #18619).
* Assignment of `None` values are no longer special, and we don't use
partial None if the feature is enabled for simple variables.
* Lvalues other than simple variables (e.g. `self.x`) continue to work
as in the past. Attribute types can't be widened, since they are
externally visible and widening could cause confusion, but this is
something we might relax in the future. Globals can be widened, however.
This seems necessary for consistency.
* If a loop body widens a variable type, we have to analyze the body
again. However, we only do one extra pass, since the inferred type could
be expanded without bound (consider `x = 0` outside loop and `x = [x]`
within the loop body).
* We first infer the type of an rvalue without using the lvalue type as
context, as otherwise the type context would often prevent redefinition.
If the rvalue type isn't valid for inference (e.g. list item type can't
be inferred), we fall back to the lvalue type context.

There are some other known bugs and limitations:
* Annotated variables can't be freely redefined (but they can still be
narrowed, of course). I may want to relax this in the future, but I'm
not sure yet.
* If there is a function definition between assignments to a variable,
the inferred types may be incorrect.
* There are few tests for `nonlocal` and some other features. We don't
have good test coverage for deferrals, mypy daemon, and disabling strict
optional.
* Imported names can't be redefined in a consistent way. This needs
further analysis.

In self check the feature generates 6 additional errors, which all seem
correct -- we infer more precise types, which will generate additional
errors due to invariant containers and fixing false negatives.

When type checking the largest internal codebase at Dropbox, this
generated about 700 new errors, the vast majority of which seemed
legitimate. Mostly they were due to inferring more precise types for
variables that used to have `Any` types. I used a recent but not the
latest version of the feature to type check the internal codebase.

Author: JukkaL

mypy/binder.py

@@ -7,8 +7,9 @@
 from typing_extensions import TypeAlias as _TypeAlias
 
 from mypy.erasetype import remove_instance_last_known_values
-from mypy.literals import Key, literal, literal_hash, subkeys
+from mypy.literals import Key, extract_var_from_literal_hash, literal, literal_hash, subkeys
 from mypy.nodes import Expression, IndexExpr, MemberExpr, NameExpr, RefExpr, TypeInfo, Var
+from mypy.options import Options
 from mypy.subtypes import is_same_type, is_subtype
 from mypy.typeops import make_simplified_union
 from mypy.types import (
@@ -39,6 +40,7 @@ class CurrentType(NamedTuple):
 
 class Frame:
     """A Frame represents a specific point in the execution of a program.
+
     It carries information about the current types of expressions at
     that point, arising either from assignments to those expressions
     or the result of isinstance checks and other type narrowing
@@ -97,7 +99,7 @@ class A:
     # This maps an expression to a list of bound types for every item in the union type.
     type_assignments: Assigns | None = None
 
-    def __init__(self) -> None:
+    def __init__(self, options: Options) -> None:
         # Each frame gets an increasing, distinct id.
         self.next_id = 1
 
@@ -131,6 +133,11 @@ def __init__(self) -> None:
         self.break_frames: list[int] = []
         self.continue_frames: list[int] = []
 
+        # If True, initial assignment to a simple variable (e.g. "x", but not "x.y")
+        # is added to the binder. This allows more precise narrowing and more
+        # flexible inference of variable types (--allow-redefinition-new).
+        self.bind_all = options.allow_redefinition_new
+
     def _get_id(self) -> int:
         self.next_id += 1
         return self.next_id
@@ -226,12 +233,20 @@ def update_from_options(self, frames: list[Frame]) -> bool:
         for key in keys:
             current_value = self._get(key)
             resulting_values = [f.types.get(key, current_value) for f in frames]
-            if any(x is None for x in resulting_values):
+            # Keys can be narrowed using two different semantics. The new semantics
+            # is enabled for plain variables when bind_all is true, and it allows
+            # variable types to be widened using subsequent assignments. This is
+            # tricky to support for instance attributes (primarily due to deferrals),
+            # so we don't use it for them.
+            old_semantics = not self.bind_all or extract_var_from_literal_hash(key) is None
+            if old_semantics and any(x is None for x in resulting_values):
                 # We didn't know anything about key before
                 # (current_value must be None), and we still don't
                 # know anything about key in at least one possible frame.
                 continue
 
+            resulting_values = [x for x in resulting_values if x is not None]
+
             if all_reachable and all(
                 x is not None and not x.from_assignment for x in resulting_values
             ):
@@ -278,7 +293,11 @@ def update_from_options(self, frames: list[Frame]) -> bool:
                     # still equivalent to such type).
                     if isinstance(type, UnionType):
                         type = collapse_variadic_union(type)
-                    if isinstance(type, ProperType) and isinstance(type, UnionType):
+                    if (
+                        old_semantics
+                        and isinstance(type, ProperType)
+                        and isinstance(type, UnionType)
+                    ):
                         # Simplify away any extra Any's that were added to the declared
                         # type when popping a frame.
                         simplified = UnionType.make_union(

mypy/build.py

@@ -2240,8 +2240,10 @@ def semantic_analysis_pass1(self) -> None:
         # TODO: Do this while constructing the AST?
         self.tree.names = SymbolTable()
         if not self.tree.is_stub:
-            # Always perform some low-key variable renaming
-            self.tree.accept(LimitedVariableRenameVisitor())
+            if not self.options.allow_redefinition_new:
+                # Perform some low-key variable renaming when assignments can't
+                # widen inferred types
+                self.tree.accept(LimitedVariableRenameVisitor())
             if options.allow_redefinition:
                 # Perform more renaming across the AST to allow variable redefinitions
                 self.tree.accept(VariableRenameVisitor())