Skip to content

Commit de6427b

Browse files
AWhetterAWhetter
committed
Can understand functools.total_ordering
1 parent 952ca23 commit de6427b

File tree

4 files changed

+219
-0
lines changed

4 files changed

+219
-0
lines changed

‎mypy/plugins/default.py

+4
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@ def get_class_decorator_hook(self, fullname: str
9393
) -> Optional[Callable[[ClassDefContext], None]]:
9494
from mypy.plugins import attrs
9595
from mypy.plugins import dataclasses
96+
from mypy.plugins import functools
9697

9798
if fullname in attrs.attr_class_makers:
9899
return attrs.attr_class_maker_callback
@@ -103,6 +104,9 @@ def get_class_decorator_hook(self, fullname: str
103104
)
104105
elif fullname in dataclasses.dataclass_makers:
105106
return dataclasses.dataclass_class_maker_callback
107+
elif fullname in functools.functools_total_ordering_makers:
108+
return functools.functools_total_ordering_maker_callback
109+
106110
return None
107111

108112

‎mypy/plugins/functools.py

+105
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
"""Plugin for supporting the functools standard library module."""
2+
from typing import Dict, NamedTuple, Optional
3+
4+
import mypy.plugin
5+
from mypy.nodes import ARG_OPT, ARG_POS, ARG_STAR2, Argument, FuncItem, Var
6+
from mypy.plugins.common import add_method_to_class
7+
from mypy.types import AnyType, CallableType, get_proper_type, Type, TypeOfAny, UnboundType
8+
9+
10+
functools_total_ordering_makers = {
11+
'functools.total_ordering',
12+
}
13+
14+
_ORDERING_METHODS = {
15+
'__lt__',
16+
'__le__',
17+
'__gt__',
18+
'__ge__',
19+
}
20+
21+
22+
_MethodInfo = NamedTuple('_MethodInfo', [('is_static', bool), ('type', CallableType)])
23+
24+
25+
def functools_total_ordering_maker_callback(ctx: mypy.plugin.ClassDefContext,
26+
auto_attribs_default: bool = False) -> None:
27+
"""Add dunder methods to classes decorated with functools.total_ordering."""
28+
if ctx.api.options.python_version < (3,):
29+
ctx.api.fail('"functools.total_ordering" is not supported in Python 2', ctx.reason)
30+
return
31+
32+
comparison_methods = _analyze_class(ctx)
33+
if not comparison_methods:
34+
ctx.api.fail(
35+
'No ordering operation defined when using "functools.total_ordering": < > <= >=',
36+
ctx.reason)
37+
return
38+
39+
# prefer __lt__ to __le__ to __gt__ to __ge__
40+
root = max(comparison_methods, key=lambda k: (comparison_methods[k] is None, k))
41+
root_method = comparison_methods[root]
42+
if not root_method:
43+
# None of the defined comparison methods can be analysed
44+
return
45+
46+
other_type = _find_other_type(root_method)
47+
bool_type = ctx.api.named_type('__builtins__.bool')
48+
ret_type = bool_type # type: Type
49+
if root_method.type.ret_type != ctx.api.named_type('__builtins__.bool'):
50+
proper_ret_type = get_proper_type(root_method.type.ret_type)
51+
if not (isinstance(proper_ret_type, UnboundType)
52+
and proper_ret_type.name.split('.')[-1] == 'bool'):
53+
ret_type = AnyType(TypeOfAny.implementation_artifact)
54+
for additional_op in _ORDERING_METHODS:
55+
# Either the method is not implemented
56+
# or has an unknown signature that we can now extrapolate.
57+
if not comparison_methods.get(additional_op):
58+
args = [Argument(Var('other', other_type), other_type, None, ARG_POS)]
59+
add_method_to_class(ctx.api, ctx.cls, additional_op, args, ret_type)
60+
61+
62+
def _find_other_type(method: _MethodInfo) -> Type:
63+
"""Find the type of the ``other`` argument in a comparison method."""
64+
first_arg_pos = 0 if method.is_static else 1
65+
cur_pos_arg = 0
66+
other_arg = None
67+
for arg_kind, arg_type in zip(method.type.arg_kinds, method.type.arg_types):
68+
if arg_kind in (ARG_POS, ARG_OPT):
69+
if cur_pos_arg == first_arg_pos:
70+
other_arg = arg_type
71+
break
72+
73+
cur_pos_arg += 1
74+
elif arg_kind != ARG_STAR2:
75+
other_arg = arg_type
76+
break
77+
78+
if other_arg is None:
79+
return AnyType(TypeOfAny.implementation_artifact)
80+
81+
return other_arg
82+
83+
84+
def _analyze_class(ctx: mypy.plugin.ClassDefContext) -> Dict[str, Optional[_MethodInfo]]:
85+
"""Analyze the class body, its parents, and return the comparison methods found."""
86+
# Traverse the MRO and collect ordering methods.
87+
comparison_methods = {} # type: Dict[str, Optional[_MethodInfo]]
88+
# Skip object because total_ordering does not use methods from object
89+
for cls in ctx.cls.info.mro[:-1]:
90+
for name in _ORDERING_METHODS:
91+
if name in cls.names and name not in comparison_methods:
92+
node = cls.names[name].node
93+
if isinstance(node, FuncItem) and isinstance(node.type, CallableType):
94+
comparison_methods[name] = _MethodInfo(node.is_static, node.type)
95+
continue
96+
97+
if isinstance(node, Var):
98+
proper_type = get_proper_type(node.type)
99+
if isinstance(proper_type, CallableType):
100+
comparison_methods[name] = _MethodInfo(node.is_staticmethod, proper_type)
101+
continue
102+
103+
comparison_methods[name] = None
104+
105+
return comparison_methods

‎mypy/test/testcheck.py

+1
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@
9191
'check-errorcodes.test',
9292
'check-annotated.test',
9393
'check-parameter-specification.test',
94+
'check-functools.test',
9495
]
9596

9697
# Tests that use Python 3.8-only AST features (like expression-scoped ignores):

‎test-data/unit/check-functools.test

+109
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
[case testTotalOrderingEqLt]
2+
from functools import total_ordering
3+
4+
@total_ordering
5+
class Ord:
6+
def __eq__(self, other: object) -> bool:
7+
return False
8+
9+
def __lt__(self, other: "Ord") -> bool:
10+
return False
11+
12+
reveal_type(Ord() < Ord()) # N: Revealed type is 'builtins.bool'
13+
reveal_type(Ord() <= Ord()) # N: Revealed type is 'builtins.bool'
14+
reveal_type(Ord() == Ord()) # N: Revealed type is 'builtins.bool'
15+
reveal_type(Ord() > Ord()) # N: Revealed type is 'builtins.bool'
16+
reveal_type(Ord() >= Ord()) # N: Revealed type is 'builtins.bool'
17+
18+
Ord() < 1 # E: Unsupported operand types for < ("Ord" and "int")
19+
Ord() <= 1 # E: Unsupported operand types for <= ("Ord" and "int")
20+
Ord() == 1
21+
Ord() > 1 # E: Unsupported operand types for > ("Ord" and "int")
22+
Ord() >= 1 # E: Unsupported operand types for >= ("Ord" and "int")
23+
[builtins fixtures/ops.pyi]
24+
[builtins fixtures/dict.pyi]
25+
26+
[case testTotalOrderingLambda]
27+
from functools import total_ordering
28+
from typing import Any, Callable
29+
30+
@total_ordering
31+
class Ord:
32+
__eq__: Callable[[Any, object], bool] = lambda self, other: False
33+
__lt__: Callable[[Any, "Ord"], bool] = lambda self, other: False
34+
35+
reveal_type(Ord() < Ord()) # N: Revealed type is 'builtins.bool'
36+
reveal_type(Ord() <= Ord()) # N: Revealed type is 'builtins.bool'
37+
reveal_type(Ord() == Ord()) # N: Revealed type is 'builtins.bool'
38+
reveal_type(Ord() > Ord()) # N: Revealed type is 'builtins.bool'
39+
reveal_type(Ord() >= Ord()) # N: Revealed type is 'builtins.bool'
40+
41+
Ord() < 1 # E: Argument 1 has incompatible type "int"; expected "Ord"
42+
Ord() <= 1 # E: Unsupported operand types for <= ("Ord" and "int")
43+
Ord() == 1
44+
Ord() > 1 # E: Unsupported operand types for > ("Ord" and "int")
45+
Ord() >= 1 # E: Unsupported operand types for >= ("Ord" and "int")
46+
[builtins fixtures/ops.pyi]
47+
[builtins fixtures/dict.pyi]
48+
49+
[case testTotalOrderingNonCallable]
50+
from functools import total_ordering
51+
52+
@total_ordering
53+
class Ord(object):
54+
def __eq__(self, other: object) -> bool:
55+
return False
56+
57+
__lt__ = 5
58+
59+
Ord() <= Ord() # E: Unsupported left operand type for <= ("Ord")
60+
Ord() > Ord() # E: "int" not callable
61+
Ord() >= Ord() # E: Unsupported left operand type for >= ("Ord")
62+
63+
[builtins fixtures/ops.pyi]
64+
[builtins fixtures/dict.pyi]
65+
66+
[case testTotalOrderingReturnNotBool]
67+
from functools import total_ordering
68+
69+
@total_ordering
70+
class Ord:
71+
def __eq__(self, other: object) -> bool:
72+
return False
73+
74+
def __lt__(self, other: "Ord") -> str:
75+
return "blah"
76+
77+
reveal_type(Ord() < Ord()) # N: Revealed type is 'builtins.str'
78+
reveal_type(Ord() <= Ord()) # N: Revealed type is 'Any'
79+
reveal_type(Ord() == Ord()) # N: Revealed type is 'builtins.bool'
80+
reveal_type(Ord() > Ord()) # N: Revealed type is 'Any'
81+
reveal_type(Ord() >= Ord()) # N: Revealed type is 'Any'
82+
83+
[builtins fixtures/ops.pyi]
84+
[builtins fixtures/dict.pyi]
85+
86+
[case testTotalOrderingAllowsAny]
87+
from functools import total_ordering
88+
89+
@total_ordering
90+
class Ord:
91+
def __eq__(self, other):
92+
return False
93+
94+
def __gt__(self, other):
95+
return False
96+
97+
reveal_type(Ord() < Ord()) # N: Revealed type is 'Any'
98+
Ord() <= Ord() # E: Unsupported left operand type for <= ("Ord")
99+
reveal_type(Ord() == Ord()) # N: Revealed type is 'Any'
100+
reveal_type(Ord() > Ord()) # N: Revealed type is 'Any'
101+
Ord() >= Ord() # E: Unsupported left operand type for >= ("Ord")
102+
103+
Ord() < 1 # E: Unsupported left operand type for < ("Ord")
104+
Ord() <= 1 # E: Unsupported left operand type for <= ("Ord")
105+
Ord() == 1
106+
Ord() > 1
107+
Ord() >= 1 # E: Unsupported left operand type for >= ("Ord")
108+
[builtins fixtures/ops.pyi]
109+
[builtins fixtures/dict.pyi]

0 commit comments

Comments
 (0)