Fix binder special-casing to apply to nested unions (#20912)

The change itself is minimal: I add a call to `flatten_nested_unions()`.
But I also do some reshuffling to only make this call _once_. Note this
change affects both `--allow-redefiniton-new` and regular mode. The
logic is a bit complicated, I double-checked it is still consistent
(note that function arguments are considered inferred for the purpose of
redefinition).

Author: ilevkivskyi

mypy/binder.py

@@ -34,6 +34,7 @@
     UnionType,
     UnpackType,
     find_unpack_in_list,
+    flatten_nested_unions,
     get_proper_type,
 )
 from mypy.typevars import fill_typevars_with_any
@@ -498,24 +499,24 @@ def assign_type(self, expr: Expression, type: Type, declared_type: Type | None)
                 # First case: a local/global variable without explicit annotation,
                 # in this case we just assign Any (essentially following the SSA logic).
                 self.put(expr, type)
-            elif isinstance(p_declared, UnionType) and any(
-                isinstance(get_proper_type(item), NoneType) for item in p_declared.items
-            ):
-                # Second case: explicit optional type, in this case we optimize for a common
-                # pattern when an untyped value used as a fallback replacing None.
-                new_items = [
-                    type if isinstance(get_proper_type(item), NoneType) else item
-                    for item in p_declared.items
-                ]
-                self.put(expr, UnionType(new_items))
-            elif isinstance(p_declared, UnionType) and any(
-                isinstance(get_proper_type(item), AnyType) for item in p_declared.items
-            ):
-                # Third case: a union already containing Any (most likely from an un-imported
-                # name), in this case we allow assigning Any as well.
-                self.put(expr, type)
+            elif isinstance(p_declared, UnionType):
+                all_items = flatten_nested_unions(p_declared.items)
+                if any(isinstance(get_proper_type(item), NoneType) for item in all_items):
+                    # Second case: explicit optional type, in this case we optimize for
+                    # a common pattern when an untyped value used as a fallback replacing None.
+                    new_items = [
+                        type if isinstance(get_proper_type(item), NoneType) else item
+                        for item in all_items
+                    ]
+                    self.put(expr, UnionType(new_items))
+                elif any(isinstance(get_proper_type(item), AnyType) for item in all_items):
+                    # Third case: a union already containing Any (most likely from
+                    # an un-imported name), in this case we allow assigning Any as well.
+                    self.put(expr, type)
+                else:
+                    # In all other cases we don't narrow to Any to minimize false negatives.
+                    self.put(expr, declared_type)
             else:
-                # In all other cases we don't narrow to Any to minimize false negatives.
                 self.put(expr, declared_type)
         elif isinstance(p_declared, AnyType):
             # Mirroring the first case above, we don't narrow to a precise type if the variable

test-data/unit/check-redefine2.test

@@ -1398,6 +1398,60 @@ def process3(items: list[str]) -> None:
     reveal_type(items)  # N: Revealed type is "builtins.list[builtins.str]"
 [builtins fixtures/primitives.pyi]
 
+[case testNewRedefineNarrowingForNestedUnionWithAny]
+# flags: --allow-redefinition-new --local-partial-types
+from typing import Any, Union
+
+a: Any
+Int = Union[int, Any]
+
+def test(x: Union[str, Int]) -> None:
+    x = a
+    reveal_type(x)  # N: Revealed type is "Any"
+    y: Union[str, Int]
+    y = a
+    reveal_type(y)  # N: Revealed type is "Any"
+
+[case testNewRedefineNarrowingForNestedUnionWithNone]
+# flags: --allow-redefinition-new --local-partial-types
+from typing import Any, Union
+
+a: Any
+Int = Union[int, None]
+
+def test(x: Union[str, Int]) -> None:
+    x = a
+    reveal_type(x)  # N: Revealed type is "Any"
+    y: Union[str, Int]
+    y = a
+    reveal_type(y)  # N: Revealed type is "builtins.str | Any | builtins.int"
+
+[case testRegularNarrowingForNestedUnionWithAny]
+from typing import Any, Union
+
+a: Any
+Int = Union[int, Any]
+
+def test(x: Union[str, Int]) -> None:
+    x = a
+    reveal_type(x)  # N: Revealed type is "Any"
+    y: Union[str, Int]
+    y = a
+    reveal_type(y)  # N: Revealed type is "Any"
+
+[case testRegularNarrowingForNestedUnionWithNone]
+from typing import Any, Union
+
+a: Any
+Int = Union[int, None]
+
+def test(x: Union[str, Int]) -> None:
+    x = a
+    reveal_type(x)  # N: Revealed type is "builtins.str | Any | builtins.int"
+    y: Union[str, Int]
+    y = a
+    reveal_type(y)  # N: Revealed type is "builtins.str | Any | builtins.int"
+
 [case testNewRedefineWidenedArgumentDeferral]
 # flags: --allow-redefinition-new --local-partial-types
 def foo(x: int) -> None: