Skip to content

Instantly share code, notes, and snippets.

@kevinastone
Last active August 29, 2015 14:09
Show Gist options
  • Save kevinastone/ebb1184659c8752c0e47 to your computer and use it in GitHub Desktop.
Save kevinastone/ebb1184659c8752c0e47 to your computer and use it in GitHub Desktop.
#!/usr/bin/env python
"""
Takes in an input string in python module format (X.Y.Z) and tries to find all
modules that start with that import path.
It takes two options.
The first is a ``recursive`` flag that allows to keep nesting looking for more modules.
The second is a regex pattern to test the results against.
>>> list(complete('logg'))
['logging']
>>> list(complete('logging.'))
['logging.config', 'logging.handlers']
>>> list(complete('logging.con'))
['logging.config']
>>> list(complete('logg', recursive=True))
['logging', 'logging.config', 'logging.handlers']
>>> list(complete('logg', pattern='andle', recursive=True))
['logging.handlers']
"""
import ast
import imp
from optparse import OptionParser
import re
import sys
SUFFIX_RE = re.compile(r'({0})$'.format('|'.join(suf[0] for suf in imp.get_suffixes())), re.I)
def strip_suffixes(name):
return SUFFIX_RE.sub('', name)
def find_modules_in_path(module_path):
import os
import os.path
if not os.path.isdir(module_path):
return
for filename in os.listdir(module_path):
full_path = os.path.join(module_path, filename)
module_name = strip_suffixes(filename)
if os.path.isfile(full_path) and module_name == filename:
# Invalid suffix, skip...
continue
elif os.path.isdir(full_path) and not os.path.exists(os.path.join(full_path, '__init__.py')):
continue
if module_name == '__init__':
# Skip __init__, we already have the base module
continue
yield module_name
def suggest_modules(module_paths, candidate):
found_matches = set()
for module_path in module_paths:
for module_name in find_modules_in_path(module_path):
# Look for matches
if module_name.startswith(candidate):
found_matches.add(module_name)
return sorted(found_matches)
def find_module(module_graph):
path = sys.path
# Iteratively search for the destination module
for module in module_graph:
path = [imp.find_module(module, path)[1]]
return path
def complete(input, pattern=None, recursive=False):
module_path = input.split('.')
parents = module_path[:-1]
candidate = module_path[-1]
try:
module = find_module(parents)
for option in suggest_modules(module, candidate):
module = '.'.join(parents + [option])
if not pattern or re.search(pattern, module):
yield module
if recursive and not module.endswith('.'):
# Also look for deeper names
for children in complete("{0}.".format(module), pattern=pattern, recursive=recursive):
yield children
except:
return
def get_symbol_names(node):
if isinstance(node, ast.ImportFrom):
for alias in node.names:
yield alias.name
elif isinstance(node, ast.Assign):
for target in node.targets:
yield target.id
elif isinstance(node, ast.ClassDef):
yield node.name
elif isinstance(node, ast.FunctionDef):
yield node.name
return
def iterate_symbols(filename):
with open(filename) as fh:
tree = ast.parse(fh.read(), filename=filename)
for node_name, nodes in ast.iter_fields(tree):
for node in nodes:
for name in get_symbol_names(node):
if not name:
continue
yield name
def complete_symbol(symbol_name, module_name, pattern=None):
"""
Return an iterator of symbol names for the given name and module.
>>> list(complete_symbol('DEFAULT_', 'logging.handlers'))
['DEFAULT_TCP_LOGGING_PORT', 'DEFAULT_UDP_LOGGING_PORT', 'DEFAULT_HTTP_LOGGING_PORT', 'DEFAULT_SOAP_LOGGING_PORT']
>>> list(complete_symbol('DEFAULT_', 'logging.handlers'), pattern='HTTP')
['DEFAULT_HTTP_LOGGING_PORT']
"""
mod_paths = find_module(module_name.split('.'))
for mod_path in mod_paths:
for name in iterate_symbols(mod_path):
if name.startswith(symbol_name):
if not pattern or re.search(pattern, name):
yield name
parser = OptionParser()
parser.add_option('-a', '--all', action='store_true', dest='all', default=False, help="find all nested modules")
parser.add_option('-m', '--module', action='store', type='string', dest='module', help="glob pattern for type of modules")
parser.add_option('-s', '--symbol', action='store', type='string', dest='symbol', help="glob pattern for the symbol name")
if __name__ == '__main__':
options, args = parser.parse_args()
input = args[0] if args else ''
if ':' in input:
module, sep, symbol = input.partition(':')
print(' '.join(sep.join((module, symbol)) for symbol in complete_symbol(symbol, module_name=module, pattern=options.symbol)))
print(' '.join(complete(input, pattern=options.module, recursive=options.all)))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment