definitions.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. #! /usr/bin/python
  2. from __future__ import print_function
  3. import argparse
  4. import ast
  5. import os
  6. import re
  7. import sys
  8. import yaml
  9. class DefinitionVisitor(ast.NodeVisitor):
  10. def __init__(self):
  11. super(DefinitionVisitor, self).__init__()
  12. self.functions = {}
  13. self.classes = {}
  14. self.names = {}
  15. self.attrs = set()
  16. self.definitions = {
  17. "def": self.functions,
  18. "class": self.classes,
  19. "names": self.names,
  20. "attrs": self.attrs,
  21. }
  22. def visit_Name(self, node):
  23. self.names.setdefault(type(node.ctx).__name__, set()).add(node.id)
  24. def visit_Attribute(self, node):
  25. self.attrs.add(node.attr)
  26. for child in ast.iter_child_nodes(node):
  27. self.visit(child)
  28. def visit_ClassDef(self, node):
  29. visitor = DefinitionVisitor()
  30. self.classes[node.name] = visitor.definitions
  31. for child in ast.iter_child_nodes(node):
  32. visitor.visit(child)
  33. def visit_FunctionDef(self, node):
  34. visitor = DefinitionVisitor()
  35. self.functions[node.name] = visitor.definitions
  36. for child in ast.iter_child_nodes(node):
  37. visitor.visit(child)
  38. def non_empty(defs):
  39. functions = {name: non_empty(f) for name, f in defs["def"].items()}
  40. classes = {name: non_empty(f) for name, f in defs["class"].items()}
  41. result = {}
  42. if functions:
  43. result["def"] = functions
  44. if classes:
  45. result["class"] = classes
  46. names = defs["names"]
  47. uses = []
  48. for name in names.get("Load", ()):
  49. if name not in names.get("Param", ()) and name not in names.get("Store", ()):
  50. uses.append(name)
  51. uses.extend(defs["attrs"])
  52. if uses:
  53. result["uses"] = uses
  54. result["names"] = names
  55. result["attrs"] = defs["attrs"]
  56. return result
  57. def definitions_in_code(input_code):
  58. input_ast = ast.parse(input_code)
  59. visitor = DefinitionVisitor()
  60. visitor.visit(input_ast)
  61. definitions = non_empty(visitor.definitions)
  62. return definitions
  63. def definitions_in_file(filepath):
  64. with open(filepath) as f:
  65. return definitions_in_code(f.read())
  66. def defined_names(prefix, defs, names):
  67. for name, funcs in defs.get("def", {}).items():
  68. names.setdefault(name, {"defined": []})["defined"].append(prefix + name)
  69. defined_names(prefix + name + ".", funcs, names)
  70. for name, funcs in defs.get("class", {}).items():
  71. names.setdefault(name, {"defined": []})["defined"].append(prefix + name)
  72. defined_names(prefix + name + ".", funcs, names)
  73. def used_names(prefix, item, defs, names):
  74. for name, funcs in defs.get("def", {}).items():
  75. used_names(prefix + name + ".", name, funcs, names)
  76. for name, funcs in defs.get("class", {}).items():
  77. used_names(prefix + name + ".", name, funcs, names)
  78. path = prefix.rstrip(".")
  79. for used in defs.get("uses", ()):
  80. if used in names:
  81. if item:
  82. names[item].setdefault("uses", []).append(used)
  83. names[used].setdefault("used", {}).setdefault(item, []).append(path)
  84. if __name__ == "__main__":
  85. parser = argparse.ArgumentParser(description="Find definitions.")
  86. parser.add_argument(
  87. "--unused", action="store_true", help="Only list unused definitions"
  88. )
  89. parser.add_argument(
  90. "--ignore", action="append", metavar="REGEXP", help="Ignore a pattern"
  91. )
  92. parser.add_argument(
  93. "--pattern", action="append", metavar="REGEXP", help="Search for a pattern"
  94. )
  95. parser.add_argument(
  96. "directories",
  97. nargs="+",
  98. metavar="DIR",
  99. help="Directories to search for definitions",
  100. )
  101. parser.add_argument(
  102. "--referrers",
  103. default=0,
  104. type=int,
  105. help="Include referrers up to the given depth",
  106. )
  107. parser.add_argument(
  108. "--referred",
  109. default=0,
  110. type=int,
  111. help="Include referred down to the given depth",
  112. )
  113. parser.add_argument(
  114. "--format", default="yaml", help="Output format, one of 'yaml' or 'dot'"
  115. )
  116. args = parser.parse_args()
  117. definitions = {}
  118. for directory in args.directories:
  119. for root, dirs, files in os.walk(directory):
  120. for filename in files:
  121. if filename.endswith(".py"):
  122. filepath = os.path.join(root, filename)
  123. definitions[filepath] = definitions_in_file(filepath)
  124. names = {}
  125. for filepath, defs in definitions.items():
  126. defined_names(filepath + ":", defs, names)
  127. for filepath, defs in definitions.items():
  128. used_names(filepath + ":", None, defs, names)
  129. patterns = [re.compile(pattern) for pattern in args.pattern or ()]
  130. ignore = [re.compile(pattern) for pattern in args.ignore or ()]
  131. result = {}
  132. for name, definition in names.items():
  133. if patterns and not any(pattern.match(name) for pattern in patterns):
  134. continue
  135. if ignore and any(pattern.match(name) for pattern in ignore):
  136. continue
  137. if args.unused and definition.get("used"):
  138. continue
  139. result[name] = definition
  140. referrer_depth = args.referrers
  141. referrers = set()
  142. while referrer_depth:
  143. referrer_depth -= 1
  144. for entry in result.values():
  145. for used_by in entry.get("used", ()):
  146. referrers.add(used_by)
  147. for name, definition in names.items():
  148. if name not in referrers:
  149. continue
  150. if ignore and any(pattern.match(name) for pattern in ignore):
  151. continue
  152. result[name] = definition
  153. referred_depth = args.referred
  154. referred = set()
  155. while referred_depth:
  156. referred_depth -= 1
  157. for entry in result.values():
  158. for uses in entry.get("uses", ()):
  159. referred.add(uses)
  160. for name, definition in names.items():
  161. if name not in referred:
  162. continue
  163. if ignore and any(pattern.match(name) for pattern in ignore):
  164. continue
  165. result[name] = definition
  166. if args.format == "yaml":
  167. yaml.dump(result, sys.stdout, default_flow_style=False)
  168. elif args.format == "dot":
  169. print("digraph {")
  170. for name, entry in result.items():
  171. print(name)
  172. for used_by in entry.get("used", ()):
  173. if used_by in result:
  174. print(used_by, "->", name)
  175. print("}")
  176. else:
  177. raise ValueError("Unknown format %r" % (args.format))