#!/usr/bin/env python # -*- coding: UTF-8 -*- """ Transform Python files into normalized import statements for grepping. """ import compiler from compiler.visitor import ASTVisitor, walk from cStringIO import StringIO import os import shlex import sys import grin __version__ = '1.2' def normalize_From(node): """ Return a list of strings of Python 'from' statements, one import on each line. """ statements = [] children = node.getChildren() module = children[0] for name, asname in children[1]: line = 'from %s import %s' % (module, name) if asname is not None: line += ' as %s' % asname line += '\n' statements.append(line) return statements def normalize_Import(node): """ Return a list of strings of Python 'import' statements, one import on each line. """ statements = [] children = node.getChildren() for name, asname in children[0]: line = 'import %s' % (name) if asname is not None: line += ' as %s' % asname line += '\n' statements.append(line) return statements class ImportPuller(ASTVisitor): """ Extract import statements from an AST. """ def __init__(self): ASTVisitor.__init__(self) self.statements = [] def visitFrom(self, node): self.statements.extend(normalize_From(node)) def visitImport(self, node): self.statements.extend(normalize_Import(node)) def as_string(self): """ Concatenate all of the 'import' and 'from' statements. """ return ''.join(self.statements) def normalize_file(filename, *args): """ Import-normalize a file. If the file is not parseable, an empty filelike object will be returned. """ try: ast = compiler.parseFile(filename) except Exception, e: return StringIO('') ip = ImportPuller() walk(ast, ip) return StringIO(ip.as_string()) def get_grinimports_arg_parser(parser=None): """ Create the command-line parser. """ parser = grin.get_grin_arg_parser(parser) parser.set_defaults(include='*.py') parser.description = ("Extract, normalize and search import statements " "from Python files.") parser.epilog = """ For example, if I have a file example.py with a bunch of imports: $ cat example.py import foo import foo.baz as blah from foo import bar, baz as bat def somefunction(): "Do something to foo.baz" import from_inside.function We can grep for 'import' in order to get all of the import statements: $ grinimports.py import example.py example.py: 1 : import foo 2 : import foo.baz as blah 3 : from foo import bar 4 : from foo import baz as bat 5 : import from_inside.function If we just want to find imports of foo.baz, we can do this: $ grinimports.py "import foo\.baz|from foo import baz" example.py example.py: 2 : import foo.baz as blah 4 : from foo import baz as bat A typical grep (or grin) cannot find all of these in the original files because the import statements are not normalized. $ grin "foo\.baz|from foo import baz" example.py example.py: 2 : import foo.baz as blah 6 : "Do something to foo.baz" """ for action in parser._actions: if hasattr(action, 'version'): action.version = 'grinpython %s' % __version__ return parser def grinimports_main(argv=None): if argv is None: # Look at the GRIN_ARGS environment variable for more arguments. env_args = shlex.split(os.getenv('GRIN_ARGS', '')) argv = [sys.argv[0]] + env_args + sys.argv[1:] parser = get_grinimports_arg_parser() args = parser.parse_args(argv[1:]) if args.context is not None: args.before_context = args.context args.after_context = args.context args.use_color = args.force_color or (not args.no_color and sys.stdout.isatty() and (os.environ.get('TERM') != 'dumb')) regex = grin.get_regex(args) g = grin.GrepText(regex, args) for filename, kind in grin.get_filenames(args): if kind == 'text': # Ignore gzipped files. report = g.grep_a_file(filename, opener=normalize_file) sys.stdout.write(report) if __name__ == '__main__': grinimports_main()