Skip to content

Commit

Permalink
Allow NotImplemented to be returned from special methods that suppo…
Browse files Browse the repository at this point in the history
…rt it (#479)
  • Loading branch information
JelleZijlstra authored Feb 12, 2022
1 parent 61dc846 commit df91898
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 1 deletion.
2 changes: 2 additions & 0 deletions docs/changelog.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
31 changes: 30 additions & 1 deletion pyanalyze/name_check_visitor.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from argparse import ArgumentParser
import ast
import enum
import itertools
from ast_decompiler import decompile
import asyncio
import builtins
Expand Down Expand Up @@ -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__"),
Expand Down Expand Up @@ -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(
Expand All @@ -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)

Expand Down
13 changes: 13 additions & 0 deletions pyanalyze/test_name_check_visitor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit df91898

Please sign in to comment.