Filter directory tree using given regex pattern.

tree-grep.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""`tree-grep.py [-i] DIR PATTERN` is similar to `tree DIR | grep [-i] PATTERN`

but it shows the whole branch upto the matched child node.


# Installation

* termcolor - helper functions for ANSII color formatting

        $ pip install termcolor

* colorama - multiplatform support (Windows)

        $ pip install colorama

* [optional] regex - new regex module is used if installed

        $ pip install regex
"""
# TBD use less memory, stream results
import os
import sys
import warnings
from functools import partial

try:
    import regex as re
except ImportError:
    warnings.warn("can't import regex module, re module will be used instead",
                  ImportWarning)
    import re

from colorama import init
init(strip=not sys.stdout.isatty())  # strip colors if stdout is redirected
from termcolor import colored


ASCII_INDENTS = "|--", "`--", "|  ", "   "


def make_tree(dirpath, parent=None):
    """Create a directory tree with dirpath as a root."""
    root = dict(parent=parent, name=dirpath, children=[])
    dirs = [(dirpath, root)]
    while dirs:
        dirpath, tree = dirs.pop()
        try:
            dir_entries = os.listdir(dirpath)
        except OSError:
            pass  # ignore errors
        else:
            for name in dir_entries:
                path = os.path.join(dirpath, name)
                entry = dict(parent=tree, name=name)
                tree['children'].append(entry)
                if os.path.isdir(path):
                    entry['children'] = []
                    dirs.append((path, entry))  # schedule to add children
    return root


def walk(tree, visit):
    """Walk the tree and call `visit` on each node in the tree."""
    stack = [tree]
    while stack:
        tree = stack.pop()
        visit(tree)
        stack += tree.get('children') or []


def mark_found(match_name, node):
    """Mark node if it matched.

    Save the matched span. Mark node parents.
    """
    matched = match_name(node['name'])
    if matched:
        # mark node
        node['marked'] = matched.span()  # save indexes within
                                         # name where the match
                                         # occured

        # mark parents
        parent = node['parent']
        while parent is not None and 'marked' not in parent:
            parent['marked'] = (1,)    # mark
            parent = parent['parent']  # go upwards


def print_marked(tree, indent=None,
                 indentation_marks=ASCII_INDENTS):
    """Print only marked nodes in the tree."""
    span = tree.get('marked')
    if not span:
        return  # skip unmark nodes

    if indent is None:
        indent = []

    # print node
    name = tree['name']
    if len(span) == 2:  # highlight the match
        print("%s%s%s%s" % (
                    "".join(indent),
                    colored(name[:span[0]], attrs=['reverse']),
                    colored(name[span[0]:span[1]], on_color='on_red',
                            attrs=['bold','reverse']),
                    colored(name[span[1]:], attrs=['reverse'])))
    else:  # no match itself but one of descendants matched
        print("%s%s" % ("".join(indent), name))

    # filter out unmarked children
    nodes = [node for node in tree.get('children') or [] if node.get('marked')]

    # print children
    if nodes:  # has marked children
        child_mark, last_child_mark, parent_mark, blank_mark = indentation_marks
        if indent:
            if indent[-1] is last_child_mark:
                indent[-1] = blank_mark
            elif indent[-1] is child_mark:
                indent[-1] = parent_mark

        # print intermediate children
        indent.append(child_mark)
        for i in xrange(len(nodes) - 1):
            print_marked(nodes[i], indent, indentation_marks)

        # print last child
        indent[-1] = last_child_mark
        print_marked(nodes[-1], indent)
        indent.pop()
        if indent and indent[-1] is parent_mark:
            indent[-1] = child_mark


def main():
    """Entry-point for the script."""
    # ignore case if `-i` is passed
    try:
        sys.argv.remove('-i')
    except ValueError:
        flags = 0
    else:
        flags = re.IGNORECASE

    if len(sys.argv) != 3:
        sys.exit(__doc__.splitlines()[0])  # print usage

    tree = make_tree(sys.argv[1])
    walk(tree, partial(mark_found, re.compile(sys.argv[2], flags=flags).search))
    print_marked(tree, indentation_marks=ASCII_INDENTS)

main()