docstring_grep.py

docstring_grep.py

#!/usr/bin/env python2.7

import ast
import itertools
import fnmatch
import os
import sys


class DocstringScanner(object):
  """Walks python files under the cwd looking for classes, functions or methods with docstrings containing substring matches."""

  @staticmethod
  def _iter_subdir_python_files(cwd='.', glob='*.py'):
    for root, dirs_to_recurse, files in os.walk(cwd, topdown=True):
      # Ignore hidden dirs.
      dirs_to_recurse[:] = [d for d in dirs_to_recurse if not d[0] == '.']
      for f in fnmatch.filter(files, glob):
        yield os.path.join(root, f)

  @classmethod
  def _iter_function_nodes_from_file(cls, filename):
    seen = set()

    def _key(node):
      return node.lineno

    def _yielder(parent, node):
      if _key(node) not in seen:
        yield parent, node
        seen.add(_key(node))

    def _walker(nodes, parent=None):
      for node in nodes:
        if isinstance(node, ast.ClassDef):
          # Yield the class node and recurse into methods declared in its body with attribution to parent.
          for t in itertools.chain(_yielder(parent, node), _walker(node.body, parent=node)):
            yield t
        elif isinstance(node, ast.FunctionDef):
          for t in _yielder(parent, node):
            yield t

    with open(filename) as f:
      parsed = ast.parse(f.read())

    for t in _walker(ast.walk(parsed)):
      yield t

  @classmethod
  def _filter_docstrings_from_file(cls, filename, substr): # (filename, function_node)
    for parent, node in cls._iter_function_nodes_from_file(filename):
      docstring = ast.get_docstring(node)
      if docstring and substr in docstring:
        yield filename, parent, node

  @classmethod
  def find(cls, substr):
    for filename in cls._iter_subdir_python_files():
      try:
        for result in cls._filter_docstrings_from_file(filename, substr):
          yield result
      except SyntaxError:
        # Seen for some intentionally syntactically broken test files.
        sys.stderr.write('[WARN] SyntaxError caught for {}'.format(filename))


def main():
  substr = sys.argv[1]
  for filename, parent, obj in DocstringScanner.find(substr):
    print('{}{}::{}'.format(filename, '::{}'.format(parent.name) if parent else '', obj.name))


if __name__ == '__main__':
  main()