Get rid of Optional when assigning Any in an is None branch

mypy/binder.py

@@ -6,7 +6,7 @@
 if MYPY:
     from typing import DefaultDict
 
-from mypy.types import Type, AnyType, PartialType, UnionType, TypeOfAny
+from mypy.types import Type, AnyType, PartialType, UnionType, TypeOfAny, NoneTyp
 from mypy.subtypes import is_subtype
 from mypy.join import join_simple
 from mypy.sametypes import is_same_type
@@ -276,6 +276,20 @@ def assign_type(self, expr: Expression,
                 # isinstance check), but now it is reassigned, so broaden back
                 # to Any (which is the most recent enclosing type)
                 self.put(expr, enclosing_type)
+        # As a special case, when assigning Any to a variable with a
+        # declared Optional type that has been narrowed to None,
+        # replace all the Nones in the declared Union type with Any.
+        # This overrides the normal behavior of ignoring Any assignments to variables
+        # in order to prevent false positives.
+        # (See discussion in #3526)
+        elif (isinstance(type, AnyType)
+              and isinstance(declared_type, UnionType)
+              and any(isinstance(item, NoneTyp) for item in declared_type.items)
+              and isinstance(self.most_recent_enclosing_type(expr, NoneTyp()), NoneTyp)):
+            # Replace any Nones in the union type with Any
+            new_items = [type if isinstance(item, NoneTyp) else item
+                         for item in declared_type.items]
+            self.put(expr, UnionType(new_items))
         elif (isinstance(type, AnyType)
               and not (isinstance(declared_type, UnionType)
                        and any(isinstance(item, AnyType) for item in declared_type.items))):

test-data/unit/check-optional.test

@@ -698,3 +698,52 @@ def g(x: str) -> None: pass
 
 y = int()
 [builtins fixtures/bool.pyi]
+
+[case testOptionalAssignAny1]
+from typing import Optional
+def f():
+    return 0
+
+def g(x: Optional[int]) -> int:
+    if x is None:
+        reveal_type(x)  # E: Revealed type is 'None'
+        # As a special case for Unions containing None, during
+        x = f()
+        reveal_type(x)  # E: Revealed type is 'Union[builtins.int, Any]'
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.int, Any]'
+    return x
+
+[builtins fixtures/bool.pyi]
+
+[case testOptionalAssignAny2]
+from typing import Optional
+def f():
+    return 0
+
+def g(x: Optional[int]) -> int:
+    if x is None:
+        reveal_type(x)  # E: Revealed type is 'None'
+        x = 1
+        reveal_type(x)  # E: Revealed type is 'builtins.int'
+        # Since we've assigned to x, the special case None behavior shouldn't happen
+        x = f()
+        reveal_type(x)  # E: Revealed type is 'Union[builtins.int, None]'
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.int, None]'
+    return x  # E: Incompatible return value type (got "Optional[int]", expected "int")
+
+[builtins fixtures/bool.pyi]
+
+[case testOptionalAssignAny3]
+from typing import Optional
+def f():
+    return 0
+
+def g(x: Optional[int]) -> int:
+    if x is not None:
+        return x
+    reveal_type(x)  # E: Revealed type is 'None'
+    x = f()
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.int, Any]'
+    return x
+
+[builtins fixtures/bool.pyi]