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