1.1 --- /dev/null Thu Jan 01 00:00:00 1970 +0000
1.2 +++ b/astgrep.py Sun Oct 26 21:02:17 2008 +0100
1.3 @@ -0,0 +1,176 @@
1.4 +#!/usr/bin/env python
1.5 +
1.6 +"""
1.7 +Search Python abstract syntax trees for nodes of a particular type having a
1.8 +particular textual value.
1.9 +
1.10 +Copyright (C) 2008 Paul Boddie <paul@boddie.org.uk>
1.11 +
1.12 +This program is free software; you can redistribute it and/or modify it under
1.13 +the terms of the GNU General Public License as published by the Free Software
1.14 +Foundation; either version 3 of the License, or (at your option) any later
1.15 +version.
1.16 +
1.17 +This program is distributed in the hope that it will be useful, but WITHOUT
1.18 +ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
1.19 +FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
1.20 +details.
1.21 +
1.22 +You should have received a copy of the GNU General Public License along with
1.23 +this program. If not, see <http://www.gnu.org/licenses/>.
1.24 +"""
1.25 +
1.26 +import compiler
1.27 +import os
1.28 +
1.29 +def search_recursive(directory, term_type, term, op=None):
1.30 +
1.31 + """
1.32 + Search files within the filesystem below 'directory' for terms having the
1.33 + given 'term_type' whose value matches the specified 'term'.
1.34 + """
1.35 +
1.36 + results = []
1.37 + for path, directories, filenames in os.walk(args["directory"]):
1.38 + for filename in filenames:
1.39 + if os.path.splitext(filename)[-1] == os.path.extsep + "py":
1.40 + results += search_file(os.path.join(path, filename), term_type, term, op)
1.41 + return results
1.42 +
1.43 +def search_file(filename, term_type, term, op=None):
1.44 +
1.45 + """
1.46 + Search the file with the given 'filename' for terms having the given
1.47 + 'term_type' whose value matches the specified 'term'.
1.48 + """
1.49 +
1.50 + node = compiler.parseFile(filename)
1.51 + cls = getattr(compiler.ast, term_type)
1.52 + return search_tree(node, cls, term, op, filename)
1.53 +
1.54 +def search_tree(node, cls, term, op=None, filename=None):
1.55 +
1.56 + """
1.57 + Search the tree rooted at the given 'node' for nodes of the given class
1.58 + 'cls' for content matching the specified 'term'.
1.59 +
1.60 + Return a list of results of the form (node, value, filename).
1.61 + """
1.62 +
1.63 + results = []
1.64 +
1.65 + if isinstance(node, cls):
1.66 + if op is None:
1.67 + results.append((node, None, filename))
1.68 + else:
1.69 + for child in node.getChildren():
1.70 + if isinstance(child, (str, unicode, int, float, long, bool)) and op(unicode(child), term):
1.71 + results.append((node, child, filename))
1.72 + break
1.73 +
1.74 + # Search within nodes, even if matches have already been found.
1.75 +
1.76 + for child in node.getChildNodes():
1.77 + results += search_tree(child, cls, term, op, filename)
1.78 +
1.79 + return results
1.80 +
1.81 +def get_operator(name):
1.82 +
1.83 + "Return the operator having the given 'name'."
1.84 +
1.85 + if name is None:
1.86 + return unicode.__eq__
1.87 + else:
1.88 + return getattr(unicode, name)
1.89 +
1.90 +def expand_results(results):
1.91 +
1.92 + """
1.93 + Expand the given 'results', making a list containing tuples of the form
1.94 + (filename, line number, line, value).
1.95 + """
1.96 +
1.97 + expanded = []
1.98 +
1.99 + for node, value, filename in results:
1.100 + if filename is not None:
1.101 + line = linecache.getline(filename, node.lineno).rstrip()
1.102 + else:
1.103 + line = None
1.104 +
1.105 + expanded.append((filename, node.lineno, line, value))
1.106 +
1.107 + return expanded
1.108 +
1.109 +# Command syntax.
1.110 +
1.111 +syntax_description = """
1.112 + [ -n ]
1.113 + [ -v ]
1.114 + -t <term-type>
1.115 + [ -o <operator> ]
1.116 + [ -e <expression> ]
1.117 + ( ( -r <directory> ) | <filename> )
1.118 + """
1.119 +
1.120 +# Main program.
1.121 +
1.122 +if __name__ == "__main__":
1.123 + import sys
1.124 + import cmdsyntax
1.125 + import linecache
1.126 +
1.127 + # Match command arguments.
1.128 +
1.129 + syntax = cmdsyntax.Syntax(syntax_description)
1.130 + syntax_matches = syntax.get_args(sys.argv[1:])
1.131 +
1.132 + try:
1.133 + args = syntax_matches[0]
1.134 + except IndexError:
1.135 + print "Syntax:"
1.136 + print syntax_description
1.137 + sys.exit(1)
1.138 +
1.139 + # Get the search details.
1.140 +
1.141 + term_type = args["term-type"]
1.142 + term = args.get("expression")
1.143 +
1.144 + if term is None:
1.145 + op = None
1.146 + else:
1.147 + op = get_operator(args.get("operator"))
1.148 +
1.149 + # Perform the search either in a single file or in a directory hierarchy.
1.150 +
1.151 + if args.has_key("filename"):
1.152 + results = search_file(args["filename"], term_type, term, op)
1.153 + else:
1.154 + results = search_recursive(args["directory"], term_type, term, op)
1.155 +
1.156 + # Present the results.
1.157 +
1.158 + for filename, lineno, line, value in expand_results(results):
1.159 + format = "%s:"
1.160 + output = [filename]
1.161 +
1.162 + if args.has_key("n"):
1.163 + format += "%d:"
1.164 + output.append(lineno)
1.165 +
1.166 + if args.has_key("v"):
1.167 + if value is not None:
1.168 + format += "%r:"
1.169 + output.append(value)
1.170 + else:
1.171 + format += "%s:"
1.172 + output.append("<%s>" % term_type)
1.173 +
1.174 + format += " %s"
1.175 + output.append(line)
1.176 +
1.177 + print format % tuple(output)
1.178 +
1.179 +# vim: tabstop=4 expandtab shiftwidth=4