definitions.py 6.4 KB

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