Fix type narrowing of == None and in (None,) conditions

mypy/checker.py

@@ -215,7 +215,13 @@
     is_literal_type,
     is_named_instance,
 )
-from mypy.types_utils import is_optional, remove_optional, store_argument_type, strip_type
+from mypy.types_utils import (
+    can_be_none,
+    is_optional,
+    remove_optional,
+    store_argument_type,
+    strip_type,
+)
 from mypy.typetraverser import TypeTraverserVisitor
 from mypy.typevars import fill_typevars, fill_typevars_with_any, has_no_typevars
 from mypy.util import is_dunder, is_sunder, is_typeshed_file
@@ -5658,7 +5664,7 @@ def has_no_custom_eq_checks(t: Type) -> bool:
                             )
                             if (
                                 collection_item_type is not None
-                                and not is_optional(collection_item_type)
+                                and not can_be_none(collection_item_type)
                                 and not (
                                     isinstance(collection_item_type, Instance)
                                     and collection_item_type.type.fullname == "builtins.object"
@@ -6065,7 +6071,7 @@ def refine_away_none_in_comparison(
         non_optional_types = []
         for i in chain_indices:
             typ = operand_types[i]
-            if not is_optional(typ):
+            if not can_be_none(typ):
                 non_optional_types.append(typ)
 
         # Make sure we have a mixture of optional and non-optional types.

mypy/plugins/common.py

@@ -43,7 +43,7 @@
     deserialize_type,
     get_proper_type,
 )
-from mypy.types_utils import is_optional
+from mypy.types_utils import can_be_none
 from mypy.typevars import fill_typevars
 from mypy.util import get_unique_redefinition_name
 
@@ -141,7 +141,7 @@ def find_shallow_matching_overload_item(overload: Overloaded, call: CallExpr) ->
                         break
                 elif (
                     arg_none
-                    and not is_optional(arg_type)
+                    and not can_be_none(arg_type)
                     and not (
                         isinstance(arg_type, Instance)
                         and arg_type.type.fullname == "builtins.object"

mypy/types_utils.py

@@ -108,6 +108,11 @@ def is_optional(t: Type) -> bool:
     )
 
 
+def can_be_none(t: Type) -> bool:
+    t = get_proper_type(t)
+    return isinstance(t, NoneType) or is_optional(t)
+
+
 def remove_optional(typ: Type) -> Type:
     typ = get_proper_type(typ)
     if isinstance(typ, UnionType):

test-data/unit/check-narrowing.test

@@ -1261,3 +1261,30 @@ def g() -> None:
         def foo(): ...
     foo()
 [builtins fixtures/dict.pyi]
+
+
+[case testNarrowingOptionalEqualsNone]
+from typing import Optional
+
+class A: ...
+
+val: Optional[A]
+
+if val == None:
+    reveal_type(val)  # N: Revealed type is "Union[__main__.A, None]"
+else:
+    reveal_type(val)  # N: Revealed type is "Union[__main__.A, None]"
+if val != None:
+    reveal_type(val)  # N: Revealed type is "Union[__main__.A, None]"
+else:
+    reveal_type(val)  # N: Revealed type is "Union[__main__.A, None]"
+
+if val in (None,):
+    reveal_type(val)  # N: Revealed type is "Union[__main__.A, None]"
+else:
+    reveal_type(val)  # N: Revealed type is "Union[__main__.A, None]"
+if val not in (None,):
+    reveal_type(val)  # N: Revealed type is "Union[__main__.A, None]"
+else:
+    reveal_type(val)  # N: Revealed type is "Union[__main__.A, None]"
+[builtins fixtures/primitives.pyi]

test-data/unit/fixtures/primitives.pyi

@@ -45,7 +45,8 @@ class memoryview(Sequence[int]):
     def __iter__(self) -> Iterator[int]: pass
     def __contains__(self, other: object) -> bool: pass
     def __getitem__(self, item: int) -> int: pass
-class tuple(Generic[T]): pass
+class tuple(Generic[T]):
+    def __contains__(self, other: object) -> bool: pass
 class list(Sequence[T]):
     def __iter__(self) -> Iterator[T]: pass
     def __contains__(self, other: object) -> bool: pass