todo.py - produce warnings from TODOs created by the current user

todo.py

from __future__ import print_function
from __future__ import unicode_literals

import argparse
import os
import re

from contextlib import closing
from itertools import chain
from itertools import count
from multiprocessing import Pool

TODO_REGEX = re.compile(r'(?:\/\/|\/\*)[^\/\*]*(TODO|FIXME)\s*\((\w+)\):?(.+)', flags=re.IGNORECASE)

FILE_EXTENSIONS = set(['.m', '.h', '.swift'])

def parse_file(path):
    user = os.getenv('USER')
    warnings = []
    with open(path, 'r') as source_file:
        for line, line_num in zip(source_file, count(1)):
            m = TODO_REGEX.search(line)
            if m and m.group(2) == user:
                warnings.append('warning: {} - {}:{}:{}'.format(m.group(1), m.group(3).strip(), path, line_num))
    return warnings

def is_good_file_ext(path):
    return os.path.splitext(path)[-1] in FILE_EXTENSIONS

def iter_paths(root):
    for dirpath, _, filenames in os.walk(root):
        for filename in filenames:
            yield os.path.join(dirpath, filename)

def main():
    parser = argparse.ArgumentParser(description='Scans for TODOs and emits warnings')
    parser.add_argument('root', type=str, help='Source root to scan')

    source_root = parser.parse_args().root

    with closing(Pool(processes=4)) as p:
        warnings = p.map(parse_file, iter_paths(source_root))

    for warning in chain.from_iterable(warnings):
        print(warning)

    return os.EX_OK

if __name__ == '__main__':
    exit(main())