Allow redefinitions for function arguments (#20853)

Fixes #19918

Implementation is straightforward, but logic behind is essentially a
compromise. Both empty context and full type context for redefinition
have pros and cons. We select full context first for function arguments
because:
* It avoids false positives in at least some of the "covariant
normalization" patterns. For example `list[str | None]` to `list[str]`
_in a branch_ will infer `list[str | None]` instead of `list[str | None]
| list[str]`. Unconditional normalization like `list[str | None]` to
`list[str]` still doesn't work, but it doesn't work for empty context
either.
* Conceptually function arguments are in between annotated and annotated
variables, in the sense that user doesn't really have an option to leave
them unannotated if they want the redefinition semantics. So we allow
redefinition but use type context behavior like for annotated variables.

cc @JukkaL

Author: ilevkivskyi

mypy/cache.py

@@ -69,7 +69,7 @@
 from mypy_extensions import u8
 
 # High-level cache layout format
-CACHE_VERSION: Final = 5
+CACHE_VERSION: Final = 6
 
 # Type used internally to represent errors:
 #   (path, line, column, end_line, end_column, severity, message, code)

mypy/checker.py

@@ -4442,7 +4442,10 @@ def check_lvalue(
             if (
                 self.options.allow_redefinition_new
                 and isinstance(lvalue.node, Var)
-                and lvalue.node.is_inferred
+                # We allow redefinition for function arguments inside function body.
+                # Although we normally do this for variables without annotation, users
+                # don't have a choice to leave a function argument without annotation.
+                and (lvalue.node.is_inferred or lvalue.node.is_argument)
             ):
                 inferred = lvalue.node
             self.store_type(lvalue, lvalue_type)
@@ -4710,8 +4713,15 @@ def infer_rvalue_with_fallback_context(
         # There are two cases where we want to try re-inferring r.h.s. in a fallback
         # type context. First case is when redefinitions are allowed, and we got
         # invalid type when using the preferred (empty) type context.
-        redefinition_fallback = inferred is not None and not is_valid_inferred_type(
-            rvalue_type, self.options
+        redefinition_fallback = (
+            inferred is not None
+            and not inferred.is_argument
+            and not is_valid_inferred_type(rvalue_type, self.options)
+        )
+        # For function arguments the preference order is opposite, and we use errors
+        # during type-checking as the fallback trigger.
+        argument_redefinition_fallback = (
+            inferred is not None and inferred.is_argument and local_errors.has_new_errors()
         )
         # Try re-inferring r.h.s. in empty context for union with explicit annotation,
         # and use it results in a narrower type. This helps with various practical
@@ -4723,7 +4733,8 @@ def infer_rvalue_with_fallback_context(
         )
 
         # Skip literal types, as they have special logic (for better errors).
-        if (redefinition_fallback or union_fallback) and not is_literal_type_like(rvalue_type):
+        try_fallback = redefinition_fallback or union_fallback or argument_redefinition_fallback
+        if try_fallback and not is_literal_type_like(rvalue_type):
             with (
                 self.msg.filter_errors(save_filtered_errors=True) as alt_local_errors,
                 self.local_type_map as alt_type_map,
@@ -4737,6 +4748,7 @@ def infer_rvalue_with_fallback_context(
                 and (
                     # For redefinition fallback we are fine getting not a subtype.
                     redefinition_fallback
+                    or argument_redefinition_fallback
                     # Skip Any type, since it is special cased in binder.
                     or not isinstance(get_proper_type(alt_rvalue_type), AnyType)
                     and is_proper_subtype(alt_rvalue_type, rvalue_type)
@@ -4773,8 +4785,8 @@ def check_simple_assignment(
             )
 
             # If redefinitions are allowed (i.e. we have --allow-redefinition-new
-            # and a variable without annotation) then we start with an empty context,
-            # since this gives somewhat more intuitive behavior. The only exception
+            # and a variable without annotation) or if a variable has union type we
+            # try inferring r.h.s. twice with a fallback type context. The only exception
             # is TypedDicts, they are often useless without context.
             try_fallback = (
                 inferred is not None or isinstance(get_proper_type(lvalue_type), UnionType)
@@ -4785,7 +4797,9 @@ def check_simple_assignment(
                     rvalue, type_context=lvalue_type, always_allow_any=always_allow_any
                 )
             else:
-                if inferred is not None:
+                # Prefer full type context for function arguments as this reduces
+                # false positives, see issue #19918 for discussion.
+                if inferred is not None and not inferred.is_argument:
                     preferred = None
                     fallback = lvalue_type
                 else:

mypy/fastparse.py

@@ -1145,6 +1145,7 @@ def make_argument(
 
         var = Var(arg.arg, arg_type)
         var.is_inferred = False
+        var.is_argument = True
         argument = Argument(var, arg_type, self.visit(default), kind, pos_only)
         argument.set_line(arg.lineno, arg.col_offset, arg.end_lineno, arg.end_col_offset)
         return argument

mypy/nodes.py

@@ -1292,6 +1292,7 @@ class Var(SymbolNode):
         "has_explicit_value",
         "allow_incompatible_override",
         "invalid_partial_type",
+        "is_argument",
     )
 
     __match_args__ = ("name", "type", "final_value")
@@ -1352,6 +1353,8 @@ def __init__(self, name: str, type: mypy.types.Type | None = None) -> None:
         # If True, this means we didn't manage to infer full type and fall back to
         # something like list[Any]. We may decide to not use such types as context.
         self.invalid_partial_type = False
+        # Is it a variable symbol for a function argument?
+        self.is_argument = False
 
     @property
     def name(self) -> str:
@@ -1415,8 +1418,6 @@ def write(self, data: WriteBuffer) -> None:
         write_flags(
             data,
             [
-                self.is_self,
-                self.is_cls,
                 self.is_initialized_in_class,
                 self.is_staticmethod,
                 self.is_classmethod,
@@ -1454,8 +1455,6 @@ def read(cls, data: ReadBuffer) -> Var:
         v.setter_type = setter_type
         v._fullname = read_str(data)
         (
-            v.is_self,
-            v.is_cls,
             v.is_initialized_in_class,
             v.is_staticmethod,
             v.is_classmethod,
@@ -1475,7 +1474,7 @@ def read(cls, data: ReadBuffer) -> Var:
             v.from_module_getattr,
             v.has_explicit_value,
             v.allow_incompatible_override,
-        ) = read_flags(data, num_flags=21)
+        ) = read_flags(data, num_flags=19)
         tag = read_tag(data)
         if tag == LITERAL_COMPLEX:
             v.final_value = complex(read_float_bare(data), read_float_bare(data))

test-data/unit/check-redefine2.test

@@ -562,8 +562,8 @@ def f2() -> None:
 def f3(x) -> None:
     if int():
         x = 0
-        reveal_type(x) # N: Revealed type is "Any"
-    reveal_type(x) # N: Revealed type is "Any"
+        reveal_type(x) # N: Revealed type is "builtins.int"
+    reveal_type(x) # N: Revealed type is "Any | builtins.int"
 
 [case tetNewRedefineDel]
 # flags: --allow-redefinition-new --local-partial-types
@@ -1310,3 +1310,32 @@ if bool():
 else:
     x = [Sub(), Sub()]
     reveal_type(x[0])  # N: Revealed type is "__main__.Sub"
+
+[case testNewRedefineFunctionArgumentsFullContext]
+# flags: --allow-redefinition-new --local-partial-types
+from typing import Optional
+
+def process(items: list[Optional[str]]) -> None:
+    if not items:
+        items = [None]
+        reveal_type(items)  # N: Revealed type is "builtins.list[builtins.str | None]"
+    process(items)  # OK
+
+[case testNewRedefineFunctionArgumentsEmptyContext]
+# flags: --allow-redefinition-new --local-partial-types
+def process1(items: list[str]) -> None:
+    items = [item.split() for item in items]
+    reveal_type(items)  # N: Revealed type is "builtins.list[builtins.list[builtins.str]]"
+
+def process2(items: list[str]) -> None:
+    if bool():
+        items = [item.split() for item in items]
+        reveal_type(items)  # N: Revealed type is "builtins.list[builtins.list[builtins.str]]"
+    reveal_type(items)  # N: Revealed type is "builtins.list[builtins.str] | builtins.list[builtins.list[builtins.str]]"
+
+def process3(items: list[str]) -> None:
+    if bool():
+        items = []
+        reveal_type(items)  # N: Revealed type is "builtins.list[builtins.str]"
+    reveal_type(items)  # N: Revealed type is "builtins.list[builtins.str]"
+[builtins fixtures/primitives.pyi]