Create Graphviz diagrams of Django templates

template_graph.py

import os
import functools
import fnmatch
import re
import collections

import pydot

def Walk(root='.', recurse=True, *patterns):
    """
        Generator for walking a directory tree.
        Starts at specified root folder, returning files
        that match our pattern. Optionally will also
        recurse through sub-folders.
    """
    for path, subdirs, files in os.walk(root):
        for name in files:
            if any(map(functools.partial(fnmatch.fnmatch, name), patterns)):
                yield os.path.join(path, name)
        if not recurse:
            break

def get_template_dirs(root_dir):
    template_dirs = [ name for name in os.listdir(root_dir)
                      if os.path.isdir(os.path.join(root_dir, name)) ]
    return template_dirs

def build_relationship_dict(directory, root_dir):
    relationship_dict = {}
    # using [^_] keeps the pattern from matching inline
    pattern = re.compile(r"""{%\s*?(include|extends|inline)\s*?['"](.*?)['"]\s*?%}""")
    for path in Walk(os.path.join(root_dir, directory), True, '*.html', '*.htm', '*.htmf', '*.ac', '*.fbml', '*.json'):
        # skip templates in the default directory
        if '/default/' not in path:
            print path
    
            file = open(path, 'r')
            text = file.read()
            file.close()
    
            matches = pattern.findall(text) # [('extends', 'messaging/base.html'), ('include', 'messaging/post.htmf')]
            # remove the templates portion of the path
            relationship_dict[path.replace('%s/' % root_dir, '')] = matches
    return relationship_dict

def build_graph(relationship_dict):
    graph = pydot.Dot(graph_type='digraph')

    # build nodes
    node_dict = {}
    for key in relationship_dict:
        a_node = pydot.Node(key, shape='box', fontname='Arial')
        graph.add_node(a_node)
        node_dict[key] = a_node

        for relationship, related_file in relationship_dict[key]:
            # add a node for the related_file if it doesn't exist yet
            if related_file not in node_dict.keys():
                a_node = pydot.Node(related_file, shape='box', fontname='Arial')
                graph.add_node(a_node)
                node_dict[related_file] = a_node

    # build edges
    for key, value in relationship_dict.iteritems():
        for relationship, related_file in value:
            if relationship == 'extends':
                graph.add_edge(pydot.Edge(node_dict[related_file], node_dict[key], color='blue'))
            if relationship == 'include':
                graph.add_edge(pydot.Edge(node_dict[key], node_dict[related_file], color='red'))
            if relationship == 'inline':
                graph.add_edge(pydot.Edge(node_dict[key], node_dict[related_file], color='green'))
    
    return graph

def main():
    output_dir = 'template_graphs'
    # make the output directory
    try:
        os.mkdir(output_dir)
    except:
        pass
    
    root_dir = 'templates'
    template_dirs = get_template_dirs(root_dir)
    for directory in template_dirs:
        relationship_dict = build_relationship_dict(directory, root_dir)
        graph = build_graph(relationship_dict)
        graph.write_png('%s/%s.png' % (output_dir, directory))

if __name__ == '__main__':
    main()