erase NotImplementedType from the return type of normally applied binary methods ("+", "and", etc.).

Author: tyralla

mypy/checker.py

@@ -4894,10 +4894,19 @@ def infer_context_dependent(
         self.store_types(original_type_map)
         return typ
 
-    def check_return_stmt(self, s: ReturnStmt) -> None:
+    @staticmethod
+    def is_notimplemented(t: ProperType) -> Type:
+        return (isinstance(t, Instance) and t.type.fullname == "builtins._NotImplementedType")
+
+    @classmethod
+    def erase_notimplemented(cls, t: ProperType) -> Type:
+        if cls.is_notimplemented(t):
+            return AnyType(TypeOfAny.special_form)
+        if isinstance(t, UnionType):
+            return UnionType.make_union([i for i in t.items if not cls.is_notimplemented(i)])
+        return t
 
-        def is_notimplemented(t: object) -> bool:
-            return isinstance(t, Instance) and t.type.fullname == "builtins._NotImplementedType"
+    def check_return_stmt(self, s: ReturnStmt) -> None:
 
         defn = self.scope.current_function()
         if defn is not None:
@@ -4979,12 +4988,7 @@ def is_notimplemented(t: object) -> bool:
                 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)]
-                            )
+                        typ_ = self.erase_notimplemented(typ)
                     self.check_subtype(
                         subtype_label="got",
                         subtype=typ_,

mypy/checkexpr.py

@@ -3554,7 +3554,7 @@ def visit_op_expr(self, e: OpExpr) -> Type:
             else:
                 assert_never(use_reverse)
             e.method_type = method_type
-            return result
+            return self.chk.erase_notimplemented(result)
         else:
             raise RuntimeError(f"Unknown operator {e.op}")
 
@@ -3705,7 +3705,7 @@ def visit_comparison_expr(self, e: ComparisonExpr) -> Type:
                 result = join.join_types(result, sub_result)
 
         assert result is not None
-        return result
+        return self.chk.erase_notimplemented(result)
 
     def find_partial_type_ref_fast_path(self, expr: Expression) -> Type | None:
         """If expression has a partial generic type, return it without additional checks.
@@ -4228,15 +4228,16 @@ def check_op(
             # callable types.
             results_final = make_simplified_union(all_results)
             inferred_final = self.combine_function_signatures(get_proper_types(all_inferred))
-            return results_final, inferred_final
+            return self.chk.erase_notimplemented(results_final), inferred_final
         else:
-            return self.check_method_call_by_name(
+            result, inferred = self.check_method_call_by_name(
                 method=method,
                 base_type=base_type,
                 args=[arg],
                 arg_kinds=[ARG_POS],
                 context=context,
             )
+            return self.chk.erase_notimplemented(result), inferred
 
     def check_boolean_op(self, e: OpExpr) -> Type:
         """Type check a boolean operation ('and' or 'or')."""

test-data/unit/check-overloading.test

@@ -6887,3 +6887,59 @@ class A:
     def i(self) -> Union[bool, NotImplementedType]: return NotImplemented
     def j(self) -> Union[bool, NotImplementedType]: return True
 [builtins fixtures/notimplemented.pyi]
+
+[case testNotImplementedReturnedFromBinaryMagicMethod]
+# flags: --warn-unreachable
+from typing import Union
+
+class A:
+    def __add__(self, x: A) -> Union[int, NotImplementedType]: ...
+    def __sub__(self, x: A) -> NotImplementedType: ...
+    def __imul__(self, x: A) -> Union[A, NotImplementedType]: ...
+    def __itruediv__(self, x: A) -> Union[A, NotImplementedType]: ...
+    def __ifloordiv__(self, x: A) -> Union[int, NotImplementedType]: ...
+    def __eq__(self, x: object) -> Union[bool, NotImplementedType]: ...
+    def __le__(self, x: int) -> Union[bool, NotImplementedType]: ...
+    def __lt__(self, x: int) -> NotImplementedType: ...
+    def __and__(self, x: object) -> NotImplementedType: ...
+class B(A):
+    def __radd__(self, x: A) -> Union[int, NotImplementedType]: ...
+    def __rsub__(self, x: A) -> NotImplementedType: ...
+    def __itruediv__(self, x: A) -> Union[A, NotImplementedType]: ...
+    def __ror__(self, x: object) -> NotImplementedType: ...
+
+a: A
+b: B
+
+reveal_type(a.__add__(a))  # N: Revealed type is "Union[builtins.int, builtins._NotImplementedType]"
+reveal_type(a.__sub__(a))  # N: Revealed type is "builtins._NotImplementedType"
+reveal_type(a.__imul__(a))  # N: Revealed type is "Union[__main__.A, builtins._NotImplementedType]"
+reveal_type(a.__eq__(a))  # N: Revealed type is "Union[builtins.bool, builtins._NotImplementedType]"
+reveal_type(a.__le__(1))  # N: Revealed type is "Union[builtins.bool, builtins._NotImplementedType]"
+
+reveal_type(a + a)  # N: Revealed type is "builtins.int"
+reveal_type(a - a)  # N: Revealed type is "Any"
+reveal_type(a + b)  # N: Revealed type is "builtins.int"
+reveal_type(a - b)  # N: Revealed type is "Any"
+def f1(a: A) -> None:
+    a += a  # E: Incompatible types in assignment (expression has type "int", variable has type "A")
+def f2(a: A) -> None:
+    a -= a
+    reveal_type(a)  # N: Revealed type is "__main__.A"
+def f3(a: A) -> None:
+    a *= a
+    reveal_type(a)  # N: Revealed type is "__main__.A"
+def f4(a: A) -> None:
+    a /= a
+    reveal_type(a)  # N: Revealed type is "__main__.A"
+def f5(a: A) -> None:
+    a //= a  # E: Result type of // incompatible in assignment
+reveal_type(a == a)  # N: Revealed type is "builtins.bool"
+reveal_type(a == 1)  # N: Revealed type is "builtins.bool"
+reveal_type(a <= 1)  # N: Revealed type is "builtins.bool"
+reveal_type(a < 1)  # N: Revealed type is "Any"
+reveal_type(a and int())  # N: Revealed type is "Union[__main__.A, builtins.int]"
+reveal_type(int() or a)  # N: Revealed type is "Union[builtins.int, __main__.A]"
+
+[builtins fixtures/notimplemented.pyi]
+

test-data/unit/fixtures/notimplemented.pyi

@@ -12,6 +12,7 @@ class str: pass
 class tuple: pass
 class dict: pass
 class classmethod: pass
+class ellipsis: pass
 
 class _NotImplementedType:
     __call__: NotImplemented  # type: ignore