py-completer.py

#!/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)))