bash/zsh-completion for python click tools

completion.py

# -*- coding: utf-8 -*-
"""
Bash completion for click command line interfaces.

Provides bash completion helpers for command-line interfaces.
"""
# batteries included
from __future__ import absolute_import

# third party
import click


def completion(cli, comp_words, comp_cword):
    """Perform bash completion."""
    words = comp_words.split()
    index = comp_cword - 1

    try:
        cur = words[index]
    except IndexError:
        cur = None

    if index == 0:
        compreply = ' '.join(sorted(cli.commands.keys()))
    elif cur is not None:
        compreply = _complete_command(cli, words, index, cur)
    else:
        compreply = ''
    return compreply


def _complete_command(cli, words, index, cur):
    cmd = words[1]

    try:
        command = cli.commands[cmd]
    except KeyError:
        compreply = ''
    else:
        compreply = _get_compreply(words, index, cur, command)
    return compreply


def _get_compreply(words, index, cur, command):
    try:
        param = _get_param(index, cur, command)
        compreply = ' '.join(sorted(param.type.choices))
    except (AttributeError, StopIteration):
        params = (p for p in command.params if _get_param_name(p) not in words)
        compreply = ' '.join(sorted(map(_get_param_name, params)))
    return compreply


def _get_param(index, cur, command):
    first_param = command.params[0]
    if index == 1 and isinstance(first_param, click.core.Argument):
        param = first_param
    else:
        pname = cur.replace('--', '').replace('-', '_')
        param = next(p for p in command.params if pname and p.name == pname)
    return param


def _get_param_name(param):
    return '--{}'.format(param.name.replace('_', '-'))