Simple shell-like completer in Python and readline

completer.py

#!/usr/bin/env python
#
# Based on the following snippets:
# - http://stackoverflow.com/questions/5637124/tab-completion-in-pythons-raw-input
# - http://pysh.sourceforge.net/
#
import os
import re
import readline
import subprocess
import traceback

RE_SPACE = re.compile('.*\s+$', re.M)

class Completer(object):
    # builtins, commands and lookup:
    # - key: command
    # - value: True if path argument
    def __init__(self):
        self.builtins = None
        self.commands = {
            'foo': True,
            'bar': False,
        }
        self.cmd_lookup = None
        self.path_lookup = None

        # try to get shell builtins
        shell = os.environ.get('SHELL')
        if shell:
            p = subprocess.Popen([shell, '-c', 'compgen -b'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            out, err = p.communicate()
            self.builtins = dict((key, True) for key in out.split())
        if not self.builtins:
            self.builtins = { 'cd': True }

    def _generate(self, prefix):
        begins_with = lambda s, p=prefix: s[:len(prefix)] == prefix
        self.cmd_lookup = {}

        # append commands in PATH
        for path in map(os.path.expanduser, os.environ.get('PATH', '').split(':')):
            if not os.path.isdir(path):
                continue

            for f in filter(begins_with, os.listdir(path)):
                self.cmd_lookup[f] = True

        # append shell builtins
        for n in filter(begins_with, self.builtins.keys()):
            self.cmd_lookup[n] = self.builtins[n]

        # append COMMANDS
        for n in filter(begins_with, self.commands.keys()):
            self.cmd_lookup[n] = self.commands[n]
        #print "\nGenerated commands for prefix %s: %s" % (`prefix`, `self.cmd_lookup`)

    def _listdir(self, root):
        "List directory 'root' appending the path separator to subdirs."
        res = []
        for name in os.listdir(root):
            path = os.path.join(root, name)
            if os.path.isdir(path):
                name += os.sep
            res.append(name)
        return res

    def _complete_path(self, path=None):
        "Perform completion of filesystem path."
        if not path:
            return self._listdir('.')
        dirname, rest = os.path.split(path)
        tmp = dirname if dirname else '.'
        res = [os.path.join(dirname, p)
                for p in self._listdir(tmp) if p.startswith(rest)]
        # more than one match, or single match which does not exist (typo)
        if len(res) > 1 or not os.path.exists(path):
            return res
        # resolved to a single directory, so return list of files below it
        if os.path.isdir(path):
            return [os.path.join(path, p) for p in self._listdir(path)]
        # exact file match terminates this completion
        return [path + ' ']

    def complete(self, text, state):
        "Generic readline completion entry point."
        try:
            buffer = readline.get_line_buffer()
            line = buffer.split()
            # account for last argument ending in a space
            if line and RE_SPACE.match(buffer):
                line.append('')
    
            # command completion
            if len(line) < 2:
                if state == 0:
                    self._generate(text)
                return ([c + ' ' for c in self.cmd_lookup.keys()] + [None])[state]
                
            # check if we should do path completion
            cmd = line[0]
            if not self.cmd_lookup:
                self._generate(cmd)
            if not cmd in self.cmd_lookup or not self.cmd_lookup[cmd]:
                return None
    
            if state == 0:
                self.path_lookup = self._complete_path(os.path.expanduser(text))
            return (self.path_lookup + [None])[state]
        except:
            traceback.print_exc()

comp = Completer()
# we want to treat '/' as part of a word, so override the delimiters
readline.set_completer_delims(' \t\n')
if 'libedit' in readline.__doc__:
    readline.parse_and_bind("bind ^I rl_complete")
else:
    readline.parse_and_bind("tab: complete")
readline.set_completer(comp.complete)
while True:
    cmd = raw_input('$ ')
    os.system(cmd)

# vi: ts=4:sw=4:et: