Skip to content

Commit 0bc2f49

Browse files
author
Guido van Rossum
committed
Fixer for 2to3 that inserts mypy annotations into all methods.
You may have to mess with the 2to3 infrastructure slightly to use this.
1 parent 640b7cf commit 0bc2f49

1 file changed

Lines changed: 216 additions & 0 deletions

File tree

‎misc/fix_annotate.py‎

Lines changed: 216 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,216 @@
1+
"""Fixer that inserts mypy annotations into all methods.
2+
3+
This transforms e.g.
4+
5+
def foo(self, bar, baz=12):
6+
return bar + baz
7+
8+
into
9+
10+
def foo(self, bar, baz=12):
11+
# type: (Any, int) -> Any
12+
return bar + baz
13+
14+
It does not do type inference but it recognizes some basic default
15+
argument values such as numbers and strings (and assumes their type
16+
implies the argument type).
17+
18+
It also uses some basic heuristics to decide whether to ignore the
19+
first argument:
20+
21+
- always if it's named 'self'
22+
- if there's a @classmethod decorator
23+
24+
Finally, it knows that __init__() is supposed to return None.
25+
"""
26+
27+
from __future__ import print_function
28+
29+
import os
30+
import re
31+
32+
from lib2to3.fixer_base import BaseFix
33+
from lib2to3.patcomp import compile_pattern
34+
from lib2to3.pytree import Leaf, Node
35+
from lib2to3.fixer_util import token, syms, touch_import
36+
37+
38+
class FixAnnotate(BaseFix):
39+
40+
# This fixer is compatible with the bottom matcher.
41+
BM_compatible = True
42+
43+
# This fixer shouldn't run by default.
44+
explicit = True
45+
46+
# The pattern to match.
47+
PATTERN = """
48+
funcdef< 'def' name=any parameters< '(' [args=any] ')' > ':' suite=any+ >
49+
"""
50+
51+
counter = None if not os.getenv('MAXFIXES') else int(os.getenv('MAXFIXES'))
52+
53+
def transform(self, node, results):
54+
if FixAnnotate.counter is not None:
55+
if FixAnnotate.counter <= 0:
56+
return
57+
suite = results['suite']
58+
children = suite[0].children
59+
60+
# NOTE: I've reverse-engineered the structure of the parse tree.
61+
# It's always a list of nodes, the first of which contains the
62+
# entire suite. Its children seem to be:
63+
#
64+
# [0] NEWLINE
65+
# [1] INDENT
66+
# [2...n-2] statements (the first may be a docstring)
67+
# [n-1] DEDENT
68+
#
69+
# Comments before the suite are part of the INDENT's prefix.
70+
#
71+
# "Compact" functions (e.g. "def foo(x, y): return max(x, y)")
72+
# have a different structure that isn't matched by PATTERN.
73+
74+
## print('-'*60)
75+
## print(node)
76+
## for i, ch in enumerate(children):
77+
## print(i, repr(ch.prefix), repr(ch))
78+
79+
# Check if there's already an annotation.
80+
for ch in children:
81+
if ch.prefix.lstrip().startswith('# type:'):
82+
return # There's already a # type: comment here; don't change anything.
83+
84+
# Compute the annotation
85+
annot = self.make_annotation(node, results)
86+
87+
# Insert '# type: {annot}' comment.
88+
# For reference, see lib2to3/fixes/fix_tuple_params.py in stdlib.
89+
if len(children) >= 2 and children[1].type == token.INDENT:
90+
children[1].prefix = '%s# type: %s\n%s' % (children[1].value, annot, children[1].prefix)
91+
children[1].changed()
92+
if FixAnnotate.counter is not None:
93+
FixAnnotate.counter -= 1
94+
95+
# Also add 'from typing import Any' at the top.
96+
if 'Any' in annot:
97+
touch_import('typing', 'Any', node)
98+
99+
def make_annotation(self, node, results):
100+
name = results['name']
101+
assert isinstance(name, Leaf), repr(name)
102+
assert name.type == token.NAME, repr(name)
103+
decorators = self.get_decorators(node)
104+
is_method = self.is_method(node)
105+
if name.value == '__init__' or not self.has_return_exprs(node):
106+
restype = 'None'
107+
else:
108+
restype = 'Any'
109+
args = results.get('args')
110+
argtypes = []
111+
if isinstance(args, Node):
112+
children = args.children
113+
elif isinstance(args, Leaf):
114+
children = [args]
115+
else:
116+
children = []
117+
# Interpret children according to the following grammar:
118+
# (('*'|'**')? NAME ['=' expr] ','?)*
119+
stars = inferred_type = ''
120+
in_default = False
121+
at_start = True
122+
for child in children:
123+
if isinstance(child, Leaf):
124+
if child.value in ('*', '**'):
125+
stars += child.value
126+
elif child.type == token.NAME and not in_default:
127+
if not is_method or not at_start or 'staticmethod' in decorators:
128+
inferred_type = 'Any'
129+
else:
130+
# Always skip the first argument if it's named 'self'.
131+
# Always skip the first argument of a class method.
132+
if child.value == 'self' or 'classmethod' in decorators:
133+
pass
134+
else:
135+
inferred_type = 'Any'
136+
elif child.value == '=':
137+
in_default = True
138+
elif in_default and child.value != ',':
139+
if child.type == token.NUMBER:
140+
if re.match(r'\d+[lL]?$', child.value):
141+
inferred_type = 'int'
142+
else:
143+
inferred_type = 'float' # TODO: complex?
144+
elif child.type == token.STRING:
145+
if child.value.startswith(('u', 'U')):
146+
inferred_type = 'unicode'
147+
else:
148+
inferred_type = 'str'
149+
elif child.type == token.NAME and child.value in ('True', 'False'):
150+
inferred_type = 'bool'
151+
elif child.value == ',':
152+
if inferred_type:
153+
argtypes.append(stars + inferred_type)
154+
# Reset
155+
stars = inferred_type = ''
156+
in_default = False
157+
at_start = False
158+
if inferred_type:
159+
argtypes.append(stars + inferred_type)
160+
return '(' + ', '.join(argtypes) + ') -> ' + restype
161+
162+
# The parse tree has a different shape when there is a single
163+
# decorator vs. when there are multiple decorators.
164+
DECORATED = "decorated< (d=decorator | decorators< dd=decorator+ >) funcdef >"
165+
decorated = compile_pattern(DECORATED)
166+
167+
def get_decorators(self, node):
168+
"""Return a list of decorators found on a function definition.
169+
170+
This is a list of strings; only simple decorators
171+
(e.g. @staticmethod) are returned.
172+
173+
If the function is undecorated or only non-simple decorators
174+
are found, return [].
175+
"""
176+
if node.parent is None:
177+
return []
178+
results = {}
179+
if not self.decorated.match(node.parent, results):
180+
return []
181+
decorators = results.get('dd') or [results['d']]
182+
decs = []
183+
for d in decorators:
184+
for child in d.children:
185+
if isinstance(child, Leaf) and child.type == token.NAME:
186+
decs.append(child.value)
187+
return decs
188+
189+
def is_method(self, node):
190+
"""Return whether the node occurs (directly) inside a class."""
191+
node = node.parent
192+
while node is not None:
193+
if node.type == syms.classdef:
194+
return True
195+
if node.type == syms.funcdef:
196+
return False
197+
node = node.parent
198+
return False
199+
200+
RETURN_EXPR = "return_stmt< 'return' any >"
201+
return_expr = compile_pattern(RETURN_EXPR)
202+
203+
def has_return_exprs(self, node):
204+
"""Traverse the tree below node looking for 'return expr'.
205+
206+
Return True if at least 'return expr' is found, False if not.
207+
(If both 'return' and 'return expr' are found, return True.)
208+
"""
209+
results = {}
210+
if self.return_expr.match(node, results):
211+
return True
212+
for child in node.children:
213+
if child.type not in (syms.funcdef, syms.classdef):
214+
if self.has_return_exprs(child):
215+
return True
216+
return False

0 commit comments

Comments
 (0)