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 import linecache 26 import types 27 28 __version__ = "0.1.1" 29 30 # Excluded AST nodes and their names. 31 32 excluded_term_types = ["Module", "Stmt"] 33 excluded_term_cls = tuple([getattr(compiler.ast, name) for name in excluded_term_types]) 34 35 # Search functions. 36 37 def search_recursive(directory, term_type, term, op=None): 38 39 """ 40 Search files within the filesystem below 'directory' for terms having the 41 given 'term_type' whose value matches the specified 'term'. 42 """ 43 44 results = [] 45 for path, directories, filenames in os.walk(directory): 46 for filename in filenames: 47 if os.path.splitext(filename)[-1] == os.path.extsep + "py": 48 results += search_file(os.path.join(path, filename), term_type, term, op) 49 return results 50 51 def search_file(filename, term_type, term, op=None): 52 53 """ 54 Search the file with the given 'filename' for terms having the given 55 'term_type' whose value matches the specified 'term'. If 'term_type' is 56 given as "*", attempt to match any term type. 57 """ 58 59 try: 60 node = compiler.parseFile(filename) 61 except SyntaxError: 62 return [] 63 64 if term_type != "*": 65 cls = getattr(compiler.ast, term_type) 66 else: 67 cls = None 68 69 return search_tree(node, cls, term, op, filename) 70 71 def search_tree(node, cls, term, op=None, filename=None): 72 73 """ 74 Search the tree rooted at the given 'node' for nodes of the given class 75 'cls' for content matching the specified 'term'. If 'cls' is None, all node 76 types will be considered for matches. 77 78 Return a list of results of the form (node, value, filename). 79 """ 80 81 results = [] 82 83 # Ignore excluded nodes. 84 85 if isinstance(node, excluded_term_cls): 86 pass 87 88 # Test permitted nodes. 89 90 elif cls is None or isinstance(node, cls): 91 if op is None: 92 results.append((node, None, filename)) 93 else: 94 for child in node.getChildren(): 95 96 # Test literals. 97 98 if isinstance(child, (str, int, float, long, bool)): 99 if op(str(child)): 100 results.append((node, child, filename)) 101 102 # Only check a single string child value since subsequent 103 # values are typically docstrings. 104 105 if isinstance(child, str): 106 break 107 108 # Search within nodes, even if matches have already been found. 109 110 for child in node.getChildNodes(): 111 results += search_tree(child, cls, term, op, filename) 112 113 return results 114 115 def expand_results(results): 116 117 """ 118 Expand the given 'results', making a list containing tuples of the form 119 (node, filename, line number, line, value). 120 """ 121 122 expanded = [] 123 124 for node, value, filename in results: 125 lineno = node.lineno 126 127 if filename is not None and lineno is not None: 128 line = linecache.getline(filename, lineno).rstrip() 129 else: 130 line = None 131 132 expanded.append((node, filename, lineno, line, value)) 133 134 return expanded 135 136 def get_term_types(): 137 138 "Return the term types supported by the module." 139 140 term_types = [] 141 142 for name in dir(compiler.ast): 143 if name in excluded_term_types: 144 continue 145 146 obj = getattr(compiler.ast, name) 147 148 if isinstance(obj, types.ClassType) and \ 149 issubclass(obj, compiler.ast.Node) and \ 150 name[0].isupper(): 151 152 term_types.append(name) 153 154 return term_types 155 156 # Command syntax. 157 158 syntax_description = """ 159 [ -n | --line-number ] 160 [ -p | --print-token ] 161 [ ( -t TERM_TYPE ) | ( --type=TERM_TYPE ) ] 162 [ ( -e PATTERN ) | ( --regexp=PATTERN ) ] 163 [ -r | -R | --recursive ] ( FILENAME ... ) 164 """ 165 166 # Main program. 167 168 def run_command(): 169 170 "The functionality of the main program." 171 172 import sys 173 import cmdsyntax 174 import re 175 import textwrap 176 177 # Match command arguments. 178 179 syntax = cmdsyntax.Syntax(syntax_description) 180 syntax_matches = syntax.get_args(sys.argv[1:]) 181 show_syntax = 0 182 183 try: 184 args = syntax_matches[0] 185 except IndexError: 186 show_syntax = 1 187 188 if show_syntax: 189 print "Syntax:" 190 print syntax_description 191 print "Term types:" 192 print "\n".join(textwrap.wrap(", ".join(get_term_types()))) 193 sys.exit(1) 194 195 # Get the search details. 196 197 term_type = args.get("TERM_TYPE", "*") 198 term = args.get("PATTERN") 199 recursive = args.has_key("r") or args.has_key("R") or args.has_key("recursive") 200 201 if term is None: 202 op = None 203 else: 204 op = re.compile(term).search 205 206 # Perform the search in files and directory hierarchies. 207 208 results = [] 209 210 for filename in args["FILENAME"]: 211 if os.path.isfile(filename): 212 results += search_file(filename, term_type, term, op) 213 elif recursive and os.path.isdir(filename): 214 results += search_recursive(filename, term_type, term, op) 215 216 # Present the results. 217 218 for node, filename, lineno, line, value in expand_results(results): 219 format = "%s:" 220 output = [filename] 221 222 # Handle line numbers and missing details. 223 224 if args.has_key("n") or args.has_key("line-number"): 225 if lineno is not None: 226 format += "%d:" 227 output.append(lineno) 228 229 # Show matching tokens, if requested. 230 231 if args.has_key("p"): 232 if value is not None: 233 format += "%r:" 234 output.append(value) 235 else: 236 format += "%s:" 237 output.append("<%s>" % (term_type or "*")) 238 239 # Show lines, if defined. 240 241 if line is not None: 242 format += " %s" 243 output.append(line) 244 245 print format % tuple(output) 246 247 if __name__ == "__main__": 248 run_command() 249 250 # vim: tabstop=4 expandtab shiftwidth=4