astgrep

Annotated astgrep.py

12:0aea3bf44a3c
2008-10-27 Paul Boddie Added tag rel-0-1 for changeset bdb3f749c9f8
paul@0 1
#!/usr/bin/env python
paul@0 2
paul@0 3
"""
paul@0 4
Search Python abstract syntax trees for nodes of a particular type having a
paul@0 5
particular textual value.
paul@0 6
paul@0 7
Copyright (C) 2008 Paul Boddie <paul@boddie.org.uk>
paul@0 8
paul@0 9
This program is free software; you can redistribute it and/or modify it under
paul@0 10
the terms of the GNU General Public License as published by the Free Software
paul@0 11
Foundation; either version 3 of the License, or (at your option) any later
paul@0 12
version.
paul@0 13
paul@0 14
This program is distributed in the hope that it will be useful, but WITHOUT
paul@0 15
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
paul@0 16
FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
paul@0 17
details.
paul@0 18
paul@0 19
You should have received a copy of the GNU General Public License along with
paul@0 20
this program.  If not, see <http://www.gnu.org/licenses/>.
paul@0 21
"""
paul@0 22
paul@0 23
import compiler
paul@0 24
import os
paul@0 25
paul@3 26
__version__ = "0.1"
paul@3 27
paul@0 28
def search_recursive(directory, term_type, term, op=None):
paul@0 29
paul@0 30
    """
paul@0 31
    Search files within the filesystem below 'directory' for terms having the
paul@0 32
    given 'term_type' whose value matches the specified 'term'.
paul@0 33
    """
paul@0 34
paul@0 35
    results = []
paul@1 36
    for path, directories, filenames in os.walk(directory):
paul@0 37
        for filename in filenames:
paul@0 38
            if os.path.splitext(filename)[-1] == os.path.extsep + "py":
paul@0 39
                results += search_file(os.path.join(path, filename), term_type, term, op)
paul@0 40
    return results
paul@0 41
paul@0 42
def search_file(filename, term_type, term, op=None):
paul@0 43
paul@0 44
    """
paul@0 45
    Search the file with the given 'filename' for terms having the given
paul@0 46
    'term_type' whose value matches the specified 'term'.
paul@0 47
    """
paul@0 48
paul@0 49
    node = compiler.parseFile(filename)
paul@0 50
    cls = getattr(compiler.ast, term_type)
paul@0 51
    return search_tree(node, cls, term, op, filename)
paul@0 52
paul@0 53
def search_tree(node, cls, term, op=None, filename=None):
paul@0 54
paul@0 55
    """
paul@0 56
    Search the tree rooted at the given 'node' for nodes of the given class
paul@0 57
    'cls' for content matching the specified 'term'.
paul@0 58
paul@0 59
    Return a list of results of the form (node, value, filename).
paul@0 60
    """
paul@0 61
paul@0 62
    results = []
paul@0 63
paul@0 64
    if isinstance(node, cls):
paul@0 65
        if op is None:
paul@0 66
            results.append((node, None, filename))
paul@0 67
        else:
paul@0 68
            for child in node.getChildren():
paul@1 69
                if isinstance(child, (str, unicode, int, float, long, bool)) and op(unicode(child)):
paul@0 70
                    results.append((node, child, filename))
paul@0 71
                    break
paul@0 72
paul@0 73
    # Search within nodes, even if matches have already been found.
paul@0 74
paul@0 75
    for child in node.getChildNodes():
paul@0 76
        results += search_tree(child, cls, term, op, filename)
paul@0 77
paul@0 78
    return results
paul@0 79
paul@0 80
def expand_results(results):
paul@0 81
paul@0 82
    """
paul@0 83
    Expand the given 'results', making a list containing tuples of the form
paul@0 84
    (filename, line number, line, value).
paul@0 85
    """
paul@0 86
paul@0 87
    expanded = []
paul@0 88
paul@0 89
    for node, value, filename in results:
paul@0 90
        if filename is not None:
paul@0 91
            line = linecache.getline(filename, node.lineno).rstrip()
paul@0 92
        else:
paul@0 93
            line = None
paul@0 94
paul@0 95
        expanded.append((filename, node.lineno, line, value))
paul@0 96
paul@0 97
    return expanded
paul@0 98
paul@0 99
# Command syntax.
paul@0 100
paul@0 101
syntax_description = """
paul@1 102
    [ -n | --line-number ]
paul@1 103
    [ -p | --print-token ]
paul@1 104
    ( ( -t TERM_TYPE ) | ( --type=TERM_TYPE ) )
paul@1 105
    [ ( -e PATTERN ) | ( --regexp=PATTERN ) ]
paul@1 106
    ( ( ( -r | -R | --recursive ) DIRECTORY ) | FILENAME )
paul@0 107
    """
paul@0 108
paul@0 109
# Main program.
paul@0 110
paul@0 111
if __name__ == "__main__":
paul@0 112
    import sys
paul@0 113
    import cmdsyntax
paul@1 114
    import re
paul@0 115
    import linecache
paul@0 116
paul@0 117
    # Match command arguments.
paul@0 118
paul@0 119
    syntax = cmdsyntax.Syntax(syntax_description)
paul@0 120
    syntax_matches = syntax.get_args(sys.argv[1:])
paul@0 121
paul@0 122
    try:
paul@0 123
        args = syntax_matches[0]
paul@0 124
    except IndexError:
paul@0 125
        print "Syntax:"
paul@0 126
        print syntax_description
paul@0 127
        sys.exit(1)
paul@0 128
paul@0 129
    # Get the search details.
paul@0 130
paul@1 131
    term_type = args["TERM_TYPE"]
paul@1 132
    term = args.get("PATTERN")
paul@0 133
paul@0 134
    if term is None:
paul@0 135
        op = None
paul@0 136
    else:
paul@1 137
        op = re.compile(term).match
paul@0 138
paul@0 139
    # Perform the search either in a single file or in a directory hierarchy.
paul@0 140
paul@1 141
    if args.has_key("FILENAME"):
paul@1 142
        results = search_file(args["FILENAME"], term_type, term, op)
paul@0 143
    else:
paul@1 144
        results = search_recursive(args["DIRECTORY"], term_type, term, op)
paul@0 145
paul@0 146
    # Present the results.
paul@0 147
paul@0 148
    for filename, lineno, line, value in expand_results(results):
paul@0 149
        format = "%s:"
paul@0 150
        output = [filename]
paul@0 151
paul@1 152
        if args.has_key("n") or args.has_key("line-number"):
paul@0 153
            format += "%d:"
paul@0 154
            output.append(lineno)
paul@0 155
paul@1 156
        if args.has_key("p"):
paul@0 157
            if value is not None:
paul@0 158
                format += "%r:"
paul@0 159
                output.append(value)
paul@0 160
            else:
paul@0 161
                format += "%s:"
paul@0 162
                output.append("<%s>" % term_type)
paul@0 163
paul@0 164
        format += " %s"
paul@0 165
        output.append(line)
paul@0 166
paul@0 167
        print format % tuple(output)
paul@0 168
paul@0 169
# vim: tabstop=4 expandtab shiftwidth=4