A simple word autocompleter meant for use in a shell

Question

Problem

I'm writing a simple shell, and I want to have the nice autocomplete feature that bash has when you partially type a word then press tab:

Right now it can work independently from a shell, but it has features that allow it to find commands on the system. Examples:

>>> table = make_lookup_table_from(["hell", "water", "help", "air", "hello", "fire", "earth"])
>>> find_matching(table, "hel")
['hell', 'hello', 'help']


>>> table = make_lookup_table_from_path()

>>> find_matching(table, "gcc-")
['gcc-ar', 'gcc-ar-8', 'gcc-ar-9', 'gcc-nm', 'gcc-nm-8', 'gcc-nm-9', 'gcc-ranlib', 'gcc-ranlib-8', 'gcc-ranlib-9', 'gcc-8', 'gcc-9']

>>> find_matching(table, "pyth")
['python3.8', 'python3.8-config', 'python3', 'python3-qr', 'python3-futurize', 'python3-pasteurize', 'python3-tor-prompt', 'python3-config', 'python3-wsdump', 'python', 'python-argcomplete-check-easy-install-script', 'python-argcomplete-check-easy-install-script3', 'python-argcomplete-tcsh', 'python-argcomplete-tcsh3', 'python-config', 'python-faraday', 'python2-config', 'python2-futurize', 'python2-pasteurize', 'python2-pbr', 'python2', 'python2.7-config', 'python2.7']

How it works:

Each word is put into a nested dictionary letter-by-letter, then terminated by a null character to mark the end of a word:

>>> make_lookup_table_from(["hell", "water", "help", "air", "hello", "fire", "earth"])
{'h': {'e': {'l': {'l': {'\x00': {}, 'o': {'\x00': {}}}, 'p': {'\x00': {}}}}}, 'w': {'a': {'t': {'e': {'r': {'\x00': {}}}}}}, 'a': {'i': {'r': {'\x00': {}}}}, 'f': {'i': {'r': {'e': {'\x00': {}}}}}, 'e': {'a': {'r': {'t': {'h': {'\x00': {}}}}}}}

To do a lookup to find matches, the tree is walked until the common sub-dictionary is found, then each word is recursively reconstructed.

Focus:

Honestly, I've been in school focusing on other things beside code, so I've gotten a little rusty. I'm using a couple less-than-ideal techniques, so any recommendations are welcome:

• The lookup function _extract_strings makes use of recursion, because this seemed like a painful problem to solve iteratively. If I'm missing an obvious alternate way, I'd appreciate any tips there.

• In the recursive function, I'm using strings to keep track of the word "so far", and passing concatenated copies to the children for them to use. I was originally using lists so I could just append without creating a new object each time, but sharing the mutable list between recurses proved to be problematic. I'm also returning only endings from the lookup function then reconstructing the full word in find_matching. This necessitates string + for every found string though, which isn't great.

These functions actually perform surprisingly fast. I was going to setup a caching, load-from-disk-on-start system to avoid needing to reconstruct the table constantly, but it's so fast that it doesn't seem worth it. As a result, both my concerns above probably fall under "premature-optimizations", but I'd still like suggestions on them or anything else, from style to other best practices.

Code:

import os
from typing import List, Iterable, Dict


_TERMINATOR = "\0"

_PATH_KEY = "PATH"
_PATH_DELIM = ":"

Table = Dict[str, "Table"]


def _get_paths() -> List[str]:
    return os.environ[_PATH_KEY].split(_PATH_DELIM)


def _find_filenames_in(paths: List[str]) -> Iterable[str]:
    return (fname
            for path in paths
            for _, _, fnames in os.walk(path)
            for fname in fnames)


def _add_string(table: Table, string: str) -> None:
    term_string = string + _TERMINATOR

    cur_level = table
    for c in term_string:
        if c not in cur_level:
            cur_level[c] = {}

        cur_level = cur_level[c]


def make_lookup_table_from(strings: Iterable[str]) -> Table:
    table = {}
    for string in strings:
        _add_string(table, string)

    return table


def make_lookup_table_from_path() -> Table:
    paths = _get_paths()
    fnames = _find_filenames_in(paths)

    return make_lookup_table_from(fnames)


def _extract_strings(table: Table) -> Iterable[str]:
    acc = []

    def rec(cur_path: str, cur_level: Table):
        for char, child in cur_level.items():
            if char == _TERMINATOR:
                acc.append(cur_path)
            else:
                rec(cur_path + char, child)

    rec("", table)

    return acc


def find_matching(table: Table, string: str) -> Iterable[str]:
    cur_level = table
    for c in string:
        try:
            cur_level = cur_level[c]
        except KeyError:
            return []

    return [string + end for end in _extract_strings(cur_level)]

Answer 1

1. You can use dict.setdefault rather than conditionally modifying the trie. This has the benefit of only looking up the key once as opposed to a maximum of 3 times with your current code.

2. I'm not a fan of your terminology:

   • table makes me think 2d array rather than a tree.
   • I'd prefer node to cur_level.
   • What does c mean?
   • Why not just call it _add_value?

def _add_value(root: Table, string: str) -> None:
    node = root
    for char in string + _TERMINATOR:
        node = node.setdefault(char, {})

3. In _extract_strings I'd move acc = [] after the function definition so the code isn't all over the place.

4. In the recursive function, I'm using strings to keep track of the word "so far", and passing concatenated copies to the children for them to use. I was originally using lists so I could just append without creating a new object each time, but sharing the mutable list between recurses proved to be problematic. I'm also returning only endings from the lookup function then reconstructing the full word in find_matching. This necessitates string + for every found string though, which isn't great.

   When getting one value your code runs in \$O(l^2)\$ where \$l\$ is the maximum length of a string. This is because each cur_path + char is an \$O(l)\$ operation and you do it \$l\$ times.

   With the current algorithm I'd suggest following an 'eh, screw it' approach and just be happy that it's sufficiently fast. As manually dealing with the stack is no fun.

   Personally I'm not a fan of acc.append, I'd instead use yield and yield from.

   def _extract_strings(table: Table) -> Iterator[str]:
       def rec(cur_path: str, cur_level: Table):
           for char, child in cur_level.items():
               if char == _TERMINATOR:
                   yield cur_path
               else:
                   yield from rec(cur_path + char, child)
       return rec("", table)
   

5. The lookup function _extract_strings makes use of recursion, because this seemed like a painful problem to solve iteratively. If I'm missing an obvious alternate way, I'd appreciate any tips there.

   Whilst manually building the stack is possible, it's not super simple. Given that the trie is unlikely to exceed Python's 1000 stack limit you can probably ignore this.
   As touched on before when building the stack we could easily build the result at the same time, changing the \$O(l^2)\$ performance to just \$O(l)\$.

   However as you should be able to see this is an abomination. I don't think anyone wants to maintain this.

   def _extract_strings(table: Table) -> Iterator[str]:
       stack = [iter(table.items())]
       stack_value = []
       while stack:
           try:
               key, value = next(stack[-1])
           except StopIteration:
               stack.pop()
               if stack_value:
                   stack_value.pop()
               continue
           if key == '\0':
               yield ''.join(stack_value)
           stack_value.append(key)
           stack.append(iter(value.items()))
   
   
   table = {
       'b': {'a': {'r': {'\0': {}}, 'z': {'\0': {}}}},
       'f': {'o': {'o': {'\0': {}}}},
   }
   for path in _extract_strings(table):
       print(path)
   

6. I'm not a fan of a lot of your empty lines. They seem random and not needed.

7. The trie would be better described in a class. A class, even with no sugar, would make your code easier to understand, as then you don't have to think "what is _add_string" and "how do I handle table".
   I've written a few tries over the years, they may help if you decide to make it a class.

   trie = Trie()
   trie.add('foo')
   trie.add('bar')
   trie.add('baz')
   # Could use the following to add instead if you need a value
   # trie['foo'] = ???
   
   key = 'ba'
   for value in trie[key]:
       print(key + value)
   

8. The plan so far will be to have a call to make_lookup_table_from_path once when the shell is loaded, save the table in a variable outside of the main loop, then do lookups on the saved table (the more I write "table", the more I'm realizing that that's probably not the right word).

   Using a Trie here is a good solution.
   If you were rebuilding the table each time you call find_matching then a simple str.startswith would probably out perform this by miles.

Overall

Your code's style seems a little odd. But otherwise it's good. I'd suggest using a class most of all.