support unions

Author: tyralla

mypy/checker.py

@@ -4895,6 +4895,10 @@ def infer_context_dependent(
         return typ
 
     def check_return_stmt(self, s: ReturnStmt) -> None:
+
+        def is_notimplemented(t: object) -> bool:
+            return isinstance(t, Instance) and t.type.fullname == "builtins._NotImplementedType"
+
         defn = self.scope.current_function()
         if defn is not None:
             if defn.is_generator:
@@ -4946,6 +4950,7 @@ def check_return_stmt(self, s: ReturnStmt) -> None:
                 if defn.is_async_generator:
                     self.fail(message_registry.RETURN_IN_ASYNC_GENERATOR, s)
                     return
+
                 # Returning a value of type Any is always fine.
                 if isinstance(typ, AnyType):
                     # (Unless you asked to be warned in that case, and the
@@ -4971,18 +4976,18 @@ def check_return_stmt(self, s: ReturnStmt) -> None:
                     if is_lambda or isinstance(typ, NoneType):
                         return
                     self.fail(message_registry.NO_RETURN_VALUE_EXPECTED, s)
-                elif (
-                    isinstance(typ, Instance)
-                    and typ.type.fullname == "builtins._NotImplementedType"
-                    and (
-                        (defn.name in BINARY_MAGIC_METHODS or defn.name == "__subclasshook__")
-                    )
-                ):
-                    return
                 else:
+                    typ_: Type = typ
+                    if defn.name in BINARY_MAGIC_METHODS or defn.name == "__subclasshook__":
+                        if is_notimplemented(typ):
+                            return
+                        if isinstance(typ, UnionType):
+                            typ_ = UnionType.make_union(
+                                [i for i in typ.items if not is_notimplemented(i)]
+                            )
                     self.check_subtype(
                         subtype_label="got",
-                        subtype=typ,
+                        subtype=typ_,
                         supertype_label="expected",
                         supertype=return_type,
                         context=s.expr,

test-data/unit/check-overloading.test

@@ -6854,8 +6854,21 @@ reveal_type(headers)  # N: Revealed type is "Union[__main__.Headers, typing.Iter
 [builtins fixtures/isinstancelist.pyi]
 
 [case testReturnNotImplementedInBinaryMagicMethods]
+from typing import Union
 class A:
-    def __eq__(self, other: object) -> bool: return NotImplemented
+    def __add__(self, other: object) -> int:
+        return NotImplemented
+    def __radd__(self, other: object) -> Union[int, NotImplementedType]:
+        return NotImplemented
+    def __sub__(self, other: object) -> Union[int, NotImplementedType]:
+        return 1
+    def __isub__(self, other: object) -> int:
+        x: Union[int, NotImplementedType]
+        return x
+    def __mul__(self, other: object) -> Union[int, NotImplementedType]:
+        x: Union[int, NotImplementedType]
+        return x
+
 [builtins fixtures/notimplemented.pyi]
 
 [case testReturnNotImplementedABCSubclassHookMethod]