Allow NotImplemented to be returned from special methods that suppo…

…rt it (#479)

docs/changelog.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Allow `NotImplemented` to be returned from special
+  methods that support it (#479)
 - Fix bug affecting type compatibility between
   generics and literals (#474)
 - Add support for `typing.Never` and `typing_extensions.Never` (#472)

pyanalyze/name_check_visitor.py

@@ -11,6 +11,7 @@
 from argparse import ArgumentParser
 import ast
 import enum
+import itertools
 from ast_decompiler import decompile
 import asyncio
 import builtins
@@ -223,6 +224,26 @@
     ast.MatMult: ("matrix multiplication", "__matmul__", "__imatmul__", "__rmatmul__"),
 }
 
+# Certain special methods are expected to return NotImplemented if they
+# can't handle a particular argument, so that the interpreter can
+# try some other call. To support thiis, such methods are allowed to
+# return NotImplemented, even if their return annotation says otherwise.
+# This is pieced together from the CPython source code, including:
+# - Methods defined as SLOT1BIN in Objects/typeobject.c
+# - Objects/abstract.c also does the binops
+# - Rich comparison in object.c and typeobject.c
+METHODS_ALLOWING_NOTIMPLEMENTED = {
+    *itertools.chain.from_iterable(BINARY_OPERATION_TO_DESCRIPTION_AND_METHOD.values()),
+    "__eq__",
+    "__ne__",
+    "__lt__",
+    "__le__",
+    "__gt__",
+    "__ge__",
+    "__length_hint__",  # Objects/abstract.c
+    "__subclasshook__",  # Modules/_abc.c
+}
+
 UNARY_OPERATION_TO_DESCRIPTION_AND_METHOD = {
     ast.Invert: ("inversion", "__invert__"),
     ast.UAdd: ("unary positive", "__pos__"),
@@ -1572,6 +1593,14 @@ def visit_FunctionDef(self, node: FunctionDefNode) -> Value:
         if not info.is_overload and not info.is_evaluated:
             self._set_name_in_scope(node.name, node, val)
 
+        if (
+            node.name in METHODS_ALLOWING_NOTIMPLEMENTED
+            and info.return_annotation is not None
+        ):
+            expected_return = info.return_annotation | KnownValue(NotImplemented)
+        else:
+            expected_return = info.return_annotation
+
         with self.asynq_checker.set_func_name(
             node.name, async_kind=info.async_kind, is_classmethod=info.is_classmethod
         ), qcore.override(self, "yield_checker", YieldChecker(self)), qcore.override(
@@ -1581,7 +1610,7 @@ def visit_FunctionDef(self, node: FunctionDefNode) -> Value:
         ), qcore.override(
             self, "current_function", potential_function
         ), qcore.override(
-            self, "expected_return_value", info.return_annotation
+            self, "expected_return_value", expected_return
         ):
             result = self._visit_function_body(info)
 

pyanalyze/test_name_check_visitor.py

@@ -473,6 +473,19 @@ def tucotuco():
             x += "a"
             assert_is_value(x, TypedValue(int))
 
+    @assert_passes()
+    def test_binop_notimplemented(self):
+        from pyanalyze.extensions import assert_type
+
+        class Capybara:
+            def __add__(self, x: str) -> bool:
+                if not isinstance(x, str):
+                    return NotImplemented
+                return len(x) > 3
+
+        def pacarana():
+            assert_type(Capybara() + "x", bool)
+
     @assert_passes()
     def test_global_sets_value(self):
         capybara = None