definitions.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. #! /usr/bin/python
  2. import ast
  3. import yaml
  4. class DefinitionVisitor(ast.NodeVisitor):
  5. def __init__(self):
  6. super(DefinitionVisitor, self).__init__()
  7. self.functions = {}
  8. self.classes = {}
  9. self.names = {}
  10. self.attrs = set()
  11. self.definitions = {
  12. 'def': self.functions,
  13. 'class': self.classes,
  14. 'names': self.names,
  15. 'attrs': self.attrs,
  16. }
  17. def visit_Name(self, node):
  18. self.names.setdefault(type(node.ctx).__name__, set()).add(node.id)
  19. def visit_Attribute(self, node):
  20. self.attrs.add(node.attr)
  21. for child in ast.iter_child_nodes(node):
  22. self.visit(child)
  23. def visit_ClassDef(self, node):
  24. visitor = DefinitionVisitor()
  25. self.classes[node.name] = visitor.definitions
  26. for child in ast.iter_child_nodes(node):
  27. visitor.visit(child)
  28. def visit_FunctionDef(self, node):
  29. visitor = DefinitionVisitor()
  30. self.functions[node.name] = visitor.definitions
  31. for child in ast.iter_child_nodes(node):
  32. visitor.visit(child)
  33. def non_empty(defs):
  34. functions = {name: non_empty(f) for name, f in defs['def'].items()}
  35. classes = {name: non_empty(f) for name, f in defs['class'].items()}
  36. result = {}
  37. if functions: result['def'] = functions
  38. if classes: result['class'] = classes
  39. names = defs['names']
  40. uses = []
  41. for name in names.get('Load', ()):
  42. if name not in names.get('Param', ()) and name not in names.get('Store', ()):
  43. uses.append(name)
  44. uses.extend(defs['attrs'])
  45. if uses: result['uses'] = uses
  46. result['names'] = names
  47. result['attrs'] = defs['attrs']
  48. return result
  49. def definitions_in_code(input_code):
  50. input_ast = ast.parse(input_code)
  51. visitor = DefinitionVisitor()
  52. visitor.visit(input_ast)
  53. definitions = non_empty(visitor.definitions)
  54. return definitions
  55. def definitions_in_file(filepath):
  56. with open(filepath) as f:
  57. return definitions_in_code(f.read())
  58. def defined_names(prefix, defs, names):
  59. for name, funcs in defs.get('def', {}).items():
  60. names.setdefault(name, {'defined': []})['defined'].append(prefix + name)
  61. defined_names(prefix + name + ".", funcs, names)
  62. for name, funcs in defs.get('class', {}).items():
  63. names.setdefault(name, {'defined': []})['defined'].append(prefix + name)
  64. defined_names(prefix + name + ".", funcs, names)
  65. def used_names(prefix, defs, names):
  66. for name, funcs in defs.get('def', {}).items():
  67. used_names(prefix + name + ".", funcs, names)
  68. for name, funcs in defs.get('class', {}).items():
  69. used_names(prefix + name + ".", funcs, names)
  70. for used in defs.get('uses', ()):
  71. if used in names:
  72. names[used].setdefault('used', []).append(prefix.rstrip('.'))
  73. if __name__ == '__main__':
  74. import sys, os, argparse, re
  75. parser = argparse.ArgumentParser(description='Find definitions.')
  76. parser.add_argument(
  77. "--unused", action="store_true", help="Only list unused definitions"
  78. )
  79. parser.add_argument(
  80. "--ignore", action="append", metavar="REGEXP", help="Ignore a pattern"
  81. )
  82. parser.add_argument(
  83. "--pattern", action="append", metavar="REGEXP",
  84. help="Search for a pattern"
  85. )
  86. parser.add_argument(
  87. "directories", nargs='+', metavar="DIR",
  88. help="Directories to search for definitions"
  89. )
  90. args = parser.parse_args()
  91. definitions = {}
  92. for directory in args.directories:
  93. for root, dirs, files in os.walk(directory):
  94. for filename in files:
  95. if filename.endswith(".py"):
  96. filepath = os.path.join(root, filename)
  97. definitions[filepath] = definitions_in_file(filepath)
  98. names = {}
  99. for filepath, defs in definitions.items():
  100. defined_names(filepath + ":", defs, names)
  101. for filepath, defs in definitions.items():
  102. used_names(filepath + ":", defs, names)
  103. patterns = [re.compile(pattern) for pattern in args.pattern or ()]
  104. ignore = [re.compile(pattern) for pattern in args.ignore or ()]
  105. result = {}
  106. for name, definition in names.items():
  107. if patterns and not any(pattern.match(name) for pattern in patterns):
  108. continue
  109. if ignore and any(pattern.match(name) for pattern in ignore):
  110. continue
  111. if args.unused and definition.get('used'):
  112. continue
  113. result[name] = definition
  114. yaml.dump(result, sys.stdout, default_flow_style=False)