fix special-case NotImplementedType in binop signatures to match runtime semantics #1129

crates/pyrefly_types/src/stdlib.rs

@@ -115,6 +115,8 @@ pub struct Stdlib {
     builtins_type: StdlibResult<ClassType>,
     /// At runtime, `types.EllipsisType` is new in 3.10, but typeshed defines it unconditionally.
     ellipsis_type: StdlibResult<ClassType>,
+    /// At runtime, `types.NotImplementedType` is new in 3.10, but typeshed defines it unconditionally.
+    not_implemented_type: StdlibResult<ClassType>,
     /// Moved from `_typeshed` to `types` in 3.10.
     none_type: StdlibResult<ClassType>,
     function_type: StdlibResult<ClassType>,
@@ -266,6 +268,7 @@ impl Stdlib {
             enumerate: lookup_generic(builtins, "enumerate", 1),
             builtins_type: lookup_concrete(builtins, "type"),
             ellipsis_type: lookup_concrete(types, "EllipsisType"),
+            not_implemented_type: lookup_concrete(types, "NotImplementedType"),
             none_type: lookup_concrete(none_location, "NoneType"),
             iterable: lookup_generic(typing, "Iterable", 1),
             iterator: lookup_generic(typing, "Iterator", 1),
@@ -384,6 +387,10 @@ impl Stdlib {
         Self::primitive(&self.ellipsis_type)
     }
 
+    pub fn not_implemented_type(&self) -> &ClassType {
+        Self::primitive(&self.not_implemented_type)
+    }
+
     pub fn none_type(&self) -> &ClassType {
         Self::primitive(&self.none_type)
     }

pyrefly/lib/alt/operators.rs

@@ -225,6 +225,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         Some(result_ty)
     }
 
+    /// Callers distribute operands over unions. The reflected dunder receives only operand
+    /// members whose forward result can be `NotImplementedType`.
     fn try_binop_calls(
         &self,
         calls: &[(&Name, &Type, &Type)],
@@ -233,6 +235,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         context: &dyn Fn() -> ErrorContext,
     ) -> Type {
         let mut first_call = None;
+        let mut successful_ret = self.heap.mk_never();
+        let not_implemented_type = self.stdlib.not_implemented_type();
         for (dunder, target, arg) in calls {
             let method_type_dunder = self.type_of_magic_dunder_attr(
                 target,
@@ -261,13 +265,35 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             // Soft errors (e.g. unknown-argument-type) must not reject an otherwise
             // valid dunder call, so gate on hard errors only.
             if !call_errors.has_hard() {
+                let ret = match ret {
+                    Type::ClassType(cls) if cls == *not_implemented_type => continue,
+                    Type::Union(mut union) => {
+                        if let Some(index) = union.members.iter().position(
+                            |member| matches!(member, Type::ClassType(cls) if cls == not_implemented_type),
+                        ) {
+                            union.members.remove(index);
+                            let ret_without_not_implemented = self.unions(union.members);
+                            errors.extend(callee_errors);
+                            successful_ret = self.union(successful_ret, ret_without_not_implemented);
+                            continue;
+                        }
+                        Type::Union(union)
+                    }
+                    ret => ret,
+                };
                 errors.extend(callee_errors);
-                return ret;
+                return if successful_ret.is_never() {
+                    ret
+                } else {
+                    self.union(successful_ret, ret)
+                };
             } else if first_call.is_none() {
                 first_call = Some((callee_errors, call_errors, ret));
             }
         }
-        if let Some((callee_errors, call_errors, ret)) = first_call {
+        if !successful_ret.is_never() {
+            successful_ret
+        } else if let Some((callee_errors, call_errors, ret)) = first_call {
             errors.extend(callee_errors);
             errors.extend(call_errors);
             ret

pyrefly/lib/test/operators.rs

@@ -706,6 +706,97 @@ assert_type(a, A)
     "#,
 );
 
+testcase!(
+    test_binop_not_implemented_type_uses_runtime_fallback,
+    r#"
+from types import NotImplementedType
+from typing import assert_type, overload
+
+class A:
+    def __add__(self, other: object) -> NotImplementedType:
+        return NotImplemented
+
+class B:
+    def __radd__(self, other: A) -> int:
+        return 1
+
+assert_type(A() + B(), int)
+
+class C:
+    def __add__(self, other: object) -> int | NotImplementedType:
+        return 1 if isinstance(other, C) else NotImplemented
+
+class D:
+    def __radd__(self, other: C) -> str:
+        return ""
+
+assert_type(C() + D(), int | str)
+
+class E:
+    pass
+
+class F:
+    def __radd__(self, other: "G") -> str:
+        return ""
+
+class G:
+    @overload
+    def __add__(self, other: E) -> int: ...
+    @overload
+    def __add__(self, other: F) -> NotImplementedType: ...
+    def __add__(self, other: E | F) -> int | NotImplementedType:
+        return 1 if isinstance(other, E) else NotImplemented
+
+def add_union(other: E | F) -> None:
+    assert_type(G() + other, int | str)
+    "#,
+);
+
+testcase!(
+    test_binop_never_does_not_use_runtime_fallback,
+    r#"
+from typing import Never, NoReturn, assert_type
+
+class A:
+    def __add__(self, other: object) -> Never:
+        raise RuntimeError
+
+class B:
+    def __radd__(self, other: A) -> str:
+        return ""
+
+class C:
+    pass
+
+class D:
+    def __add__(self, other: object) -> NoReturn:
+        raise RuntimeError
+
+assert_type(A() + B(), Never)
+assert_type(A() + C(), Never)
+assert_type(D() + C(), Never)
+    "#,
+);
+
+testcase!(
+    test_augassign_never_does_not_use_runtime_fallback,
+    r#"
+from typing import Never, assert_type
+
+class A:
+    def __iadd__(self, other: "B") -> Never:
+        raise RuntimeError
+
+class B:
+    def __radd__(self, other: A) -> str:
+        return ""
+
+def f(a: A, b: B) -> None:
+    a += b
+    assert_type(a, Never)
+    "#,
+);
+
 // We try __iadd__ and some fallback dunders. When all fail, the least confusing option is to use __iadd__.
 testcase!(
     test_iadd_error,