grep python source strings

astgrep.py

#!/usr/bin/env python
"""
grep -r for python source strings

Allows to search for strings splitted across multiple lines
"""
from __future__ import print_function

import io
import re
import os
import ast
import sys
import itertools
import functools
import multiprocessing as mp


MODE = '%r'  #  'or %s'
FMT = '%%d: %s' % MODE

PRINT_LOCK = mp.Lock()


def as_list(func):
    @functools.wraps(func)
    def _as_list(*args, **kwargs):
        return list(func(*args, **kwargs))

    return _as_list


@as_list
def grep_file(filename, regexp):
    data = io.open(filename, mode='rb').read()
    try:
        parsed = ast.parse(data, filename)
    except SyntaxError as err:
        with PRINT_LOCK:
            print('ERROR', filename, err)
            return

    for node in ast.walk(parsed):
        if isinstance(node, ast.Str) and regexp.search(node.s):
            yield filename, node


def gather(dirs):
    for directory in dirs:
        for root, _, files in os.walk(directory):
            for file in files:
                if file.endswith('.py'):
                    yield os.path.join(root, file)


def main():
    pool = mp.Pool()

    what = sys.argv[1]
    dirs = sys.argv[2:]

    print('Searching for %r in %s' % (what, ', '.join(dirs)))
    what = re.compile(what)
    files = gather(dirs)

    grep = functools.partial(grep_file, regexp=what)
    res = itertools.chain(*pool.imap_unordered(grep, files, chunksize=10))

    for filename, node in res:
        print('%s:%d %r' % (filename, node.lineno, node.s))


if __name__ == '__main__':
    main()