Last active
August 29, 2015 14:24
-
-
Save groner/86100b343460657d0a82 to your computer and use it in GitHub Desktop.
pip install astunparse; find . -name \*.py -exec python3 findclassannotations.py {} +
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import argparse | |
import ast | |
import re | |
import sys | |
from astunparse import unparse as astunparse | |
from astunparse import dump as astdump | |
def main(): | |
parser = argparse.ArgumentParser() | |
parser.add_argument('files', nargs='+', metavar='FILES') | |
args = parser.parse_args() | |
for fn in args.files: | |
with open(fn) as fh: | |
source = fh.read() | |
try: | |
tree = ast.parse(source, fn) | |
except SyntaxError as e: | |
print(fn, e, file=sys.stderr) | |
continue | |
for clsdef in AnnotatedClassScanner().visit(tree): | |
print('{}:{}: class {}'.format(fn, clsdef.lineno, clsdef.name)) | |
for funcdef, clsdef in AnnotatedMethodScanner().visit(tree): | |
print('{}:{}: {:>{indent}}def {}.{}'.format(fn, funcdef.lineno, | |
'', clsdef.name, funcdef.name, indent=funcdef.col_offset)) | |
#for funcdef, clsdef, matches in RegexScanner(r'^(\s*(\w+)\s*=\s*self\._\2\s*)$').visit(tree): | |
# print('{}:{}: {:>{indent}}def {}.{}'.format(fn, funcdef.lineno, | |
# '', clsdef.name, funcdef.name, indent=funcdef.col_offset)) | |
# for one, two in matches: | |
# print(one) | |
class IteratingNodeVisitor(ast.NodeVisitor): | |
#def visit(self, node): | |
# """Visit a node.""" | |
# method = 'visit_' + node.__class__.__name__ | |
# visitor = getattr(self, method, self.generic_visit) | |
# return visitor(node) | |
# yield from all the time | |
def generic_visit(self, node): | |
"""Called if no explicit visitor function exists for a node.""" | |
for field, value in ast.iter_fields(node): | |
if isinstance(value, list): | |
for item in value: | |
if isinstance(item, ast.AST): | |
yield from self.visit(item) | |
elif isinstance(value, ast.AST): | |
yield from self.visit(value) | |
class AnnotatedClassScanner(IteratingNodeVisitor): | |
def visit_ClassDef(self, node): | |
for dtor in node.decorator_list: | |
dtor_src = astunparse(dtor) | |
if 'annotate' in dtor_src: | |
yield node | |
yield from self.generic_visit(node) | |
class AnnotatedMethodScanner(IteratingNodeVisitor): | |
in_class = None | |
def visit_ClassDef(self, node): | |
saved_in_class, self.in_class = self.in_class, node | |
yield from self.generic_visit(node) | |
self.in_class = saved_in_class | |
def visit_FunctionDef(self, node): | |
if self.in_class: | |
is_provider = any( | |
'provider' in astunparse(dtor) | |
for dtor in self.in_class.decorator_list ) | |
for dtor in node.decorator_list: | |
dtor_src = astunparse(dtor) | |
if 'annotate' not in dtor_src: continue | |
if 'annotate.method' in dtor_src: continue | |
if is_provider and node.name in ('get', '__init__'): continue | |
#print(ast.dump(node)) | |
yield node, self.in_class | |
saved_in_class, self.in_class = self.in_class, None | |
yield from self.generic_visit(node) | |
self.in_class = saved_in_class | |
class RegexScanner(IteratingNodeVisitor): | |
in_class = None | |
def __init__(self, pattern): | |
self.pattern = re.compile(pattern, flags=re.M) | |
def visit_ClassDef(self, node): | |
saved_in_class, self.in_class = self.in_class, node | |
yield from self.generic_visit(node) | |
self.in_class = saved_in_class | |
def visit_FunctionDef(self, node): | |
if self.in_class: | |
src = astunparse(node) | |
matches = self.pattern.findall(src) | |
if matches: | |
#print(ast.dump(node)) | |
yield node, self.in_class, matches | |
saved_in_class, self.in_class = self.in_class, None | |
yield from self.generic_visit(node) | |
self.in_class = saved_in_class | |
if __name__ == '__main__': | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment