Add flag to allow more flexible variable redefinition

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.
+        self.bind_all = options.allow_redefinition_new
+
     def _get_id(self) -> int:
         self.next_id += 1
         return self.next_id
@@ -226,12 +233,15 @@ 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):
+            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
             ):

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())

mypy/checkexpr.py

@@ -1117,8 +1117,7 @@ def try_infer_partial_type(self, e: CallExpr) -> None:
             typ = self.try_infer_partial_value_type_from_call(e, callee.name, var)
             # Var may be deleted from partial_types in try_infer_partial_value_type_from_call
             if typ is not None and var in partial_types:
-                var.type = typ
-                del partial_types[var]
+                self.chk.replace_partial_type(var, typ, partial_types)
         elif isinstance(callee.expr, IndexExpr) and isinstance(callee.expr.base, RefExpr):
             # Call 'x[y].method(...)'; may infer type of 'x' if it's a partial defaultdict.
             if callee.expr.analyzed is not None:
@@ -1136,12 +1135,12 @@ def try_infer_partial_type(self, e: CallExpr) -> None:
             if value_type is not None:
                 # Infer key type.
                 key_type = self.accept(index)
-                if mypy.checker.is_valid_inferred_type(key_type):
+                if mypy.checker.is_valid_inferred_type(key_type, self.chk.options):
                     # Store inferred partial type.
                     assert partial_type.type is not None
                     typename = partial_type.type.fullname
-                    var.type = self.chk.named_generic_type(typename, [key_type, value_type])
-                    del partial_types[var]
+                    new_type = self.chk.named_generic_type(typename, [key_type, value_type])
+                    self.chk.replace_partial_type(var, new_type, partial_types)
 
     def get_partial_var(self, ref: RefExpr) -> tuple[Var, dict[Var, Context]] | None:
         var = ref.node
@@ -1176,7 +1175,7 @@ def try_infer_partial_value_type_from_call(
             and e.arg_kinds == [ARG_POS]
         ):
             item_type = self.accept(e.args[0])
-            if mypy.checker.is_valid_inferred_type(item_type):
+            if mypy.checker.is_valid_inferred_type(item_type, self.chk.options):
                 return self.chk.named_generic_type(typename, [item_type])
         elif (
             typename in self.container_args
@@ -1188,7 +1187,7 @@ def try_infer_partial_value_type_from_call(
                 arg_typename = arg_type.type.fullname
                 if arg_typename in self.container_args[typename][methodname]:
                     if all(
-                        mypy.checker.is_valid_inferred_type(item_type)
+                        mypy.checker.is_valid_inferred_type(item_type, self.chk.options)
                         for item_type in arg_type.args
                     ):
                         return self.chk.named_generic_type(typename, list(arg_type.args))
@@ -5830,7 +5829,7 @@ def visit_conditional_expr(self, e: ConditionalExpr, allow_none_return: bool = F
             else_map, e.else_expr, context=ctx, allow_none_return=allow_none_return
         )
 
-        if not mypy.checker.is_valid_inferred_type(if_type):
+        if not mypy.checker.is_valid_inferred_type(if_type, self.chk.options):
             # Analyze the right branch disregarding the left branch.
             else_type = full_context_else_type
             # we want to keep the narrowest value of else_type for union'ing the branches

mypy/main.py

@@ -93,6 +93,13 @@ def main(
         stdout, stderr, options.hide_error_codes, hide_success=bool(options.output)
     )
 
+    if options.allow_redefinition_new and not options.local_partial_types:
+        fail(
+            "error: --local-partial-types must be enabled if using --allow-redefinition-new",
+            stderr,
+            options,
+        )
+
     if options.install_types and (stdout is not sys.stdout or stderr is not sys.stderr):
         # Since --install-types performs user input, we want regular stdout and stderr.
         fail("error: --install-types not supported in this mode of running mypy", stderr, options)
@@ -856,7 +863,15 @@ def add_invertible_flag(
         "--allow-redefinition",
         default=False,
         strict_flag=False,
-        help="Allow unconditional variable redefinition with a new type",
+        help="Allow restricted, unconditional variable redefinition with a new type",
+        group=strictness_group,
+    )
+
+    add_invertible_flag(
+        "--allow-redefinition-new",
+        default=False,
+        strict_flag=False,
+        help=argparse.SUPPRESS,  # This is still very experimental
         group=strictness_group,
     )
 

mypy/nodes.py

@@ -1073,7 +1073,8 @@ def fullname(self) -> str:
         return self._fullname
 
     def __repr__(self) -> str:
-        return f"<Var {self.fullname!r} at {hex(id(self))}>"
+        name = self.fullname or self.name
+        return f"<Var {name!r} at {hex(id(self))}>"
 
     def accept(self, visitor: NodeVisitor[T]) -> T:
         return visitor.visit_var(self)
@@ -1637,11 +1638,12 @@ def accept(self, visitor: StatementVisitor[T]) -> T:
 
 
 class MatchStmt(Statement):
-    __slots__ = ("subject", "patterns", "guards", "bodies")
+    __slots__ = ("subject", "subject_dummy", "patterns", "guards", "bodies")
 
     __match_args__ = ("subject", "patterns", "guards", "bodies")
 
     subject: Expression
+    subject_dummy: NameExpr | None
     patterns: list[Pattern]
     guards: list[Expression | None]
     bodies: list[Block]
@@ -1656,6 +1658,7 @@ def __init__(
         super().__init__()
         assert len(patterns) == len(guards) == len(bodies)
         self.subject = subject
+        self.subject_dummy = None
         self.patterns = patterns
         self.guards = guards
         self.bodies = bodies

mypy/options.py

@@ -22,6 +22,7 @@ class BuildType:
 PER_MODULE_OPTIONS: Final = {
     # Please keep this list sorted
     "allow_redefinition",
+    "allow_redefinition_new",
     "allow_untyped_globals",
     "always_false",
     "always_true",
@@ -219,6 +220,10 @@ def __init__(self) -> None:
         # and the same nesting level as the initialization
         self.allow_redefinition = False
 
+        # Allow flexible variable redefinition with an arbitrary type, in different
+        # blocks and and at different nesting levels
+        self.allow_redefinition_new = False
+
         # Prohibit equality, identity, and container checks for non-overlapping types.
         # This makes 1 == '1', 1 in ['1'], and 1 is '1' errors.
         self.strict_equality = False

mypy/plugins/functools.py

@@ -276,7 +276,7 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) -
     for i, actuals in enumerate(formal_to_actual):
         if len(bound.arg_types) == len(fn_type.arg_types):
             arg_type = bound.arg_types[i]
-            if not mypy.checker.is_valid_inferred_type(arg_type):
+            if not mypy.checker.is_valid_inferred_type(arg_type, ctx.api.options):
                 arg_type = fn_type.arg_types[i]  # bit of a hack
         else:
             # TODO: I assume that bound and fn_type have the same arguments. It appears this isn't
@@ -301,7 +301,7 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) -
             partial_names.append(fn_type.arg_names[i])
 
     ret_type = bound.ret_type
-    if not mypy.checker.is_valid_inferred_type(ret_type):
+    if not mypy.checker.is_valid_inferred_type(ret_type, ctx.api.options):
         ret_type = fn_type.ret_type  # same kind of hack as above
 
     partially_applied = fn_type.copy_modified(

mypy/semanal.py

@@ -658,6 +658,13 @@ def refresh_partial(
 
     def refresh_top_level(self, file_node: MypyFile) -> None:
         """Reanalyze a stale module top-level in fine-grained incremental mode."""
+        if self.options.allow_redefinition_new and not self.options.local_partial_types:
+            n = TempNode(AnyType(TypeOfAny.special_form))
+            n.line = 1
+            n.column = 0
+            n.end_line = 1
+            n.end_column = 0
+            self.fail("--local-partial-types must be enabled if using --allow-redefinition-new", n)
         self.recurse_into_functions = False
         self.add_implicit_module_attrs(file_node)
         for d in file_node.defs:
@@ -4335,8 +4342,10 @@ def analyze_name_lvalue(
                 else:
                     lvalue.fullname = lvalue.name
                 if self.is_func_scope():
-                    if unmangle(name) == "_":
+                    if unmangle(name) == "_" and not self.options.allow_redefinition_new:
                         # Special case for assignment to local named '_': always infer 'Any'.
+                        # This isn't needed with --allow-redefinition-new, since arbitrary
+                        # types can be assigned to '_' anyway.
                         typ = AnyType(TypeOfAny.special_form)
                         self.store_declared_types(lvalue, typ)
             if is_final and self.is_final_redefinition(kind, name):

test-data/unit/check-incremental.test

@@ -6855,3 +6855,14 @@ from .lib import NT
 [builtins fixtures/tuple.pyi]
 [out]
 [out2]
+
+[case testNewRedefineAffectsCache]
+# flags: --local-partial-types --allow-redefinition-new
+# flags2: --local-partial-types
+# flags3: --local-partial-types --allow-redefinition-new
+x = 0
+if int():
+    x = ""
+[out]
+[out2]
+main:6: error: Incompatible types in assignment (expression has type "str", variable has type "int")

test-data/unit/check-python310.test

@@ -2552,3 +2552,41 @@ def f(t: T) -> None:
         case T([K() as k]):
             reveal_type(k)  # N: Revealed type is "Tuple[builtins.int, fallback=__main__.K]"
 [builtins fixtures/tuple.pyi]
+
+[case testNewRedefineMatchBasics]
+# flags: --allow-redefinition-new --local-partial-types
+
+def f1(x: int | str | list[bytes]) -> None:
+    match x:
+        case int():
+            reveal_type(x) # N: Revealed type is "builtins.int"
+        case str(y):
+            reveal_type(y) # N: Revealed type is "builtins.str"
+        case [y]:
+            reveal_type(y) # N: Revealed type is "builtins.bytes"
+    reveal_type(y) # N: Revealed type is "Union[builtins.str, builtins.bytes]"
+
+[case testNewRedefineLoopWithMatch]
+# flags: --allow-redefinition-new --local-partial-types
+
+def f1() -> None:
+    while True:
+        x = object()
+        match x:
+            case str(y):
+                pass
+            case int():
+                pass
+        if int():
+            continue
+
+def f2() -> None:
+    for x in [""]:
+        match str():
+            case "a":
+                y = ""
+            case "b":
+                y = 1
+                return
+    reveal_type(y) # N: Revealed type is "builtins.str"
+[builtins fixtures/list.pyi]