mccabe.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347
  1. """ Meager code path measurement tool.
  2. Ned Batchelder
  3. http://nedbatchelder.com/blog/200803/python_code_complexity_microtool.html
  4. MIT License.
  5. """
  6. from __future__ import with_statement
  7. import optparse
  8. import sys
  9. import tokenize
  10. from collections import defaultdict
  11. try:
  12. import ast
  13. from ast import iter_child_nodes
  14. except ImportError: # Python 2.5
  15. from flake8.util import ast, iter_child_nodes
  16. __version__ = '0.6.1'
  17. class ASTVisitor(object):
  18. """Performs a depth-first walk of the AST."""
  19. def __init__(self):
  20. self.node = None
  21. self._cache = {}
  22. def default(self, node, *args):
  23. for child in iter_child_nodes(node):
  24. self.dispatch(child, *args)
  25. def dispatch(self, node, *args):
  26. self.node = node
  27. klass = node.__class__
  28. meth = self._cache.get(klass)
  29. if meth is None:
  30. className = klass.__name__
  31. meth = getattr(self.visitor, 'visit' + className, self.default)
  32. self._cache[klass] = meth
  33. return meth(node, *args)
  34. def preorder(self, tree, visitor, *args):
  35. """Do preorder walk of tree using visitor"""
  36. self.visitor = visitor
  37. visitor.visit = self.dispatch
  38. self.dispatch(tree, *args) # XXX *args make sense?
  39. class PathNode(object):
  40. def __init__(self, name, look="circle"):
  41. self.name = name
  42. self.look = look
  43. def to_dot(self):
  44. print('node [shape=%s,label="%s"] %d;' % (
  45. self.look, self.name, self.dot_id()))
  46. def dot_id(self):
  47. return id(self)
  48. class PathGraph(object):
  49. def __init__(self, name, entity, lineno, column=0):
  50. self.name = name
  51. self.entity = entity
  52. self.lineno = lineno
  53. self.column = column
  54. self.nodes = defaultdict(list)
  55. def connect(self, n1, n2):
  56. self.nodes[n1].append(n2)
  57. # Ensure that the destination node is always counted.
  58. self.nodes[n2] = []
  59. def to_dot(self):
  60. print('subgraph {')
  61. for node in self.nodes:
  62. node.to_dot()
  63. for node, nexts in self.nodes.items():
  64. for next in nexts:
  65. print('%s -- %s;' % (node.dot_id(), next.dot_id()))
  66. print('}')
  67. def complexity(self):
  68. """ Return the McCabe complexity for the graph.
  69. V-E+2
  70. """
  71. num_edges = sum([len(n) for n in self.nodes.values()])
  72. num_nodes = len(self.nodes)
  73. return num_edges - num_nodes + 2
  74. class PathGraphingAstVisitor(ASTVisitor):
  75. """ A visitor for a parsed Abstract Syntax Tree which finds executable
  76. statements.
  77. """
  78. def __init__(self):
  79. super(PathGraphingAstVisitor, self).__init__()
  80. self.classname = ""
  81. self.graphs = {}
  82. self.reset()
  83. def reset(self):
  84. self.graph = None
  85. self.tail = None
  86. def dispatch_list(self, node_list):
  87. for node in node_list:
  88. self.dispatch(node)
  89. def visitFunctionDef(self, node):
  90. if self.classname:
  91. entity = '%s%s' % (self.classname, node.name)
  92. else:
  93. entity = node.name
  94. name = '%d:%d: %r' % (node.lineno, node.col_offset, entity)
  95. if self.graph is not None:
  96. # closure
  97. pathnode = self.appendPathNode(name)
  98. self.tail = pathnode
  99. self.dispatch_list(node.body)
  100. bottom = PathNode("", look='point')
  101. self.graph.connect(self.tail, bottom)
  102. self.graph.connect(pathnode, bottom)
  103. self.tail = bottom
  104. else:
  105. self.graph = PathGraph(name, entity, node.lineno, node.col_offset)
  106. pathnode = PathNode(name)
  107. self.tail = pathnode
  108. self.dispatch_list(node.body)
  109. self.graphs["%s%s" % (self.classname, node.name)] = self.graph
  110. self.reset()
  111. visitAsyncFunctionDef = visitFunctionDef
  112. def visitClassDef(self, node):
  113. old_classname = self.classname
  114. self.classname += node.name + "."
  115. self.dispatch_list(node.body)
  116. self.classname = old_classname
  117. def appendPathNode(self, name):
  118. if not self.tail:
  119. return
  120. pathnode = PathNode(name)
  121. self.graph.connect(self.tail, pathnode)
  122. self.tail = pathnode
  123. return pathnode
  124. def visitSimpleStatement(self, node):
  125. if node.lineno is None:
  126. lineno = 0
  127. else:
  128. lineno = node.lineno
  129. name = "Stmt %d" % lineno
  130. self.appendPathNode(name)
  131. def default(self, node, *args):
  132. if isinstance(node, ast.stmt):
  133. self.visitSimpleStatement(node)
  134. else:
  135. super(PathGraphingAstVisitor, self).default(node, *args)
  136. def visitLoop(self, node):
  137. name = "Loop %d" % node.lineno
  138. self._subgraph(node, name)
  139. visitAsyncFor = visitFor = visitWhile = visitLoop
  140. def visitIf(self, node):
  141. name = "If %d" % node.lineno
  142. self._subgraph(node, name)
  143. def _subgraph(self, node, name, extra_blocks=()):
  144. """create the subgraphs representing any `if` and `for` statements"""
  145. if self.graph is None:
  146. # global loop
  147. self.graph = PathGraph(name, name, node.lineno, node.col_offset)
  148. pathnode = PathNode(name)
  149. self._subgraph_parse(node, pathnode, extra_blocks)
  150. self.graphs["%s%s" % (self.classname, name)] = self.graph
  151. self.reset()
  152. else:
  153. pathnode = self.appendPathNode(name)
  154. self._subgraph_parse(node, pathnode, extra_blocks)
  155. def _subgraph_parse(self, node, pathnode, extra_blocks):
  156. """parse the body and any `else` block of `if` and `for` statements"""
  157. loose_ends = []
  158. self.tail = pathnode
  159. self.dispatch_list(node.body)
  160. loose_ends.append(self.tail)
  161. for extra in extra_blocks:
  162. self.tail = pathnode
  163. self.dispatch_list(extra.body)
  164. loose_ends.append(self.tail)
  165. if node.orelse:
  166. self.tail = pathnode
  167. self.dispatch_list(node.orelse)
  168. loose_ends.append(self.tail)
  169. else:
  170. loose_ends.append(pathnode)
  171. if pathnode:
  172. bottom = PathNode("", look='point')
  173. for le in loose_ends:
  174. self.graph.connect(le, bottom)
  175. self.tail = bottom
  176. def visitTryExcept(self, node):
  177. name = "TryExcept %d" % node.lineno
  178. self._subgraph(node, name, extra_blocks=node.handlers)
  179. visitTry = visitTryExcept
  180. def visitWith(self, node):
  181. name = "With %d" % node.lineno
  182. self.appendPathNode(name)
  183. self.dispatch_list(node.body)
  184. visitAsyncWith = visitWith
  185. class McCabeChecker(object):
  186. """McCabe cyclomatic complexity checker."""
  187. name = 'mccabe'
  188. version = __version__
  189. _code = 'C901'
  190. _error_tmpl = "C901 %r is too complex (%d)"
  191. max_complexity = -1
  192. def __init__(self, tree, filename):
  193. self.tree = tree
  194. @classmethod
  195. def add_options(cls, parser):
  196. flag = '--max-complexity'
  197. kwargs = {
  198. 'default': -1,
  199. 'action': 'store',
  200. 'type': 'int',
  201. 'help': 'McCabe complexity threshold',
  202. 'parse_from_config': 'True',
  203. }
  204. config_opts = getattr(parser, 'config_options', None)
  205. if isinstance(config_opts, list):
  206. # Flake8 2.x
  207. kwargs.pop('parse_from_config')
  208. parser.add_option(flag, **kwargs)
  209. parser.config_options.append('max-complexity')
  210. else:
  211. parser.add_option(flag, **kwargs)
  212. @classmethod
  213. def parse_options(cls, options):
  214. cls.max_complexity = int(options.max_complexity)
  215. def run(self):
  216. if self.max_complexity < 0:
  217. return
  218. visitor = PathGraphingAstVisitor()
  219. visitor.preorder(self.tree, visitor)
  220. for graph in visitor.graphs.values():
  221. if graph.complexity() > self.max_complexity:
  222. text = self._error_tmpl % (graph.entity, graph.complexity())
  223. yield graph.lineno, graph.column, text, type(self)
  224. def get_code_complexity(code, threshold=7, filename='stdin'):
  225. try:
  226. tree = compile(code, filename, "exec", ast.PyCF_ONLY_AST)
  227. except SyntaxError:
  228. e = sys.exc_info()[1]
  229. sys.stderr.write("Unable to parse %s: %s\n" % (filename, e))
  230. return 0
  231. complx = []
  232. McCabeChecker.max_complexity = threshold
  233. for lineno, offset, text, check in McCabeChecker(tree, filename).run():
  234. complx.append('%s:%d:1: %s' % (filename, lineno, text))
  235. if len(complx) == 0:
  236. return 0
  237. print('\n'.join(complx))
  238. return len(complx)
  239. def get_module_complexity(module_path, threshold=7):
  240. """Returns the complexity of a module"""
  241. with open(module_path, "rU") as mod:
  242. code = mod.read()
  243. return get_code_complexity(code, threshold, filename=module_path)
  244. def _read(filename):
  245. if (2, 5) < sys.version_info < (3, 0):
  246. with open(filename, 'rU') as f:
  247. return f.read()
  248. elif (3, 0) <= sys.version_info < (4, 0):
  249. """Read the source code."""
  250. try:
  251. with open(filename, 'rb') as f:
  252. (encoding, _) = tokenize.detect_encoding(f.readline)
  253. except (LookupError, SyntaxError, UnicodeError):
  254. # Fall back if file encoding is improperly declared
  255. with open(filename, encoding='latin-1') as f:
  256. return f.read()
  257. with open(filename, 'r', encoding=encoding) as f:
  258. return f.read()
  259. def main(argv=None):
  260. if argv is None:
  261. argv = sys.argv[1:]
  262. opar = optparse.OptionParser()
  263. opar.add_option("-d", "--dot", dest="dot",
  264. help="output a graphviz dot file", action="store_true")
  265. opar.add_option("-m", "--min", dest="threshold",
  266. help="minimum complexity for output", type="int",
  267. default=1)
  268. options, args = opar.parse_args(argv)
  269. code = _read(args[0])
  270. tree = compile(code, args[0], "exec", ast.PyCF_ONLY_AST)
  271. visitor = PathGraphingAstVisitor()
  272. visitor.preorder(tree, visitor)
  273. if options.dot:
  274. print('graph {')
  275. for graph in visitor.graphs.values():
  276. if (not options.threshold or
  277. graph.complexity() >= options.threshold):
  278. graph.to_dot()
  279. print('}')
  280. else:
  281. for graph in visitor.graphs.values():
  282. if graph.complexity() >= options.threshold:
  283. print(graph.name, graph.complexity())
  284. if __name__ == '__main__':
  285. main(sys.argv[1:])