bash tab completions for gkeep

gkeep_completions

#!/usr/bin/env python3

# To install (assuming you have the bash-completions package)
#   python3 -c 'import gkeep_completions; gkeep_completions.install()'
# or if you don't (this doesn't work as well, for example you have to restart your terminal and it edits your .bashrc file)
#   python3 -c 'import gkeep_completions; gkeep_completions.install_rc()'

import os, sys, shlex, subprocess

def __get_files(path):
    """
    Gets all of the files from a path where the final component of the path may only be part of a
    filename / directory name. Uses the directory components available and then the files /
    directories in that final directory must start with the given partial path.

    Hidden files are not returned in the list unless the partial filename starts with a '.'.

    Paths that begin with ~ have the user's home directory expanded.
    """
    path = os.path.expanduser(path)
    #path = os.path.expandvars(path) # TODO: expand $ in filenames?
    
    # Blank path, list everything but hidden files
    if path == '':
        return [name for name in os.listdir('.') if name[0] != '.']

    # A complete path to a directory, list all files in it except hidden
    elif os.path.isdir(path):
        return [os.path.join(path, name) for name in os.listdir(path) if name[0] != '.']

    # Only part of a filename - get the directory part of the path
    directory = os.path.dirname(path)
    if directory == '':
        # Only part of a file name with no directory in it
        return [name for name in os.listdir('.') if name.startswith(path)]
    elif os.path.isdir(directory):
        # Only part of a file name with a directory in it
        basename = os.path.basename(path)
        return [os.path.join(directory, name) for name in os.listdir(directory)
                if name.startswith(basename)]
    
    # Invalid path
    return []

def __complete_csv(path):
    """
    Complete a CSV filename partial path. This uses __get_files and then filters to only return
    directories or files that end with .csv (ignoring case). Directory names end with a '/'. File
    names end with a space ' '. All directory / file names are escaped for the bash command line.
    """
    files = __get_files(path)
    dirs = [shlex.quote(f+os.path.sep) for f in files if os.path.isdir(f)]
    csvs = [shlex.quote(f)+' ' for f in files
            if os.path.isfile(f) and os.path.splitext(f)[1].lower() == '.csv']
    return dirs + csvs

def __is_assignment_dir(path):
    """
    Return True if the path represents an assignment directory (i.e. contains base_code, tests, and
    email.txt) and False otherwsie.
    """
    return (os.path.isdir(os.path.join(path, 'base_code')) and
        os.path.isdir(os.path.join(path, 'tests')) and
        os.path.isfile(os.path.join(path, 'email.txt')))

def __complete_directory(path, assignment_dir=True):
    """
    Complete directory names from partial path. This uses __get_files and then filters to only
    return directories (which end with a /). If assignment_dir is True (the default) then assignment
    directories end with a space instead of a /. All directory names are escaped for the bash
    command line.
    """
    # Path is already an assignment directory
    if len(path) > 0 and path[-1] != '/' and __is_assignment_dir(os.path.expanduser(path)):
        return [shlex.quote(path) + ' ']
    # Get all directories
    dirs = [f for f in __get_files(path) if os.path.isdir(f)]
    if not assignment_dir:
        # Always end with just a /
        return [shlex.quote(d+os.path.sep) for d in dirs]
    # End with a space if the directory could be an assignment (contains base_code/email.txt/tests)
    return [shlex.quote(d) + ' ' if __is_assignment_dir(d) else shlex.quote(d+os.path.sep)
            for d in dirs]

def __complete_word(word, possibilities):
    """
    Return a list of the complete words from a partial word given the list of possibilities. The
    return possible words have a space added to the end of them.
    """
    return [possibility+' ' for possibility in possibilities if possibility.startswith(word)]

def __gkeep_query(query_type, max_age=15):
    """
    Return the results of a gkeep query for either assignments or students and returns a dictionary
    with the course name as the key and the values being lists of the results.

    This will create a persistent cache (even in between runs of this program) for faster access.
    The max_age of the cache in seconds can be set and defaults to 15 seconds.
    """
    import tempfile, getpass, pickle
    from datetime import datetime
    
    # Attempt to get the results from the cache
    user = getpass.getuser()
    tmp = os.path.join(tempfile.gettempdir(), 'gkeep-completions-'+user)
    file = os.path.join(tmp, query_type)
    try:
        age = datetime.now() - datetime.fromtimestamp(os.path.getmtime(file))
        if age.total_seconds() <= max_age:
            # TODO: touch?
            with open(file, 'rb') as f: return pickle.load(f)
    except OSError: pass # file doesn't exist or similar
    
    # Run the command and get the results
    info = subprocess.run(['gkeep', 'query', query_type], capture_output=True, text=True)
    if info.returncode != 0: return []
    lines = info.stdout.splitlines()

    # Create the dictionary
    data = { }
    new_class = True
    for line in lines:
        if line == '': new_class = True
        elif new_class:
            class_results = []
            data[line[:-1]] = class_results
            new_class = False
        else:
            class_results.append(line)

    # Cache the results
    os.makedirs(tmp, 0o700, True)
    with open(file, 'wb') as f: pickle.dump(data, f)

    # Return the results
    return data

def __complete_class_name(class_name):
    """
    Return a list of the complete class names (with spaces at the end) from a partial class name.
    This requires running `gkeep query classes`.
    """
    data = __gkeep_query('assignments').keys()
    words = __complete_word(class_name, data)
    if not words:
        data = __gkeep_query('assignments', 0.5).keys()
        words = __complete_word(class_name, data)
    return words

def __complete_assignment(class_name, assignment):
    """
    Return a list of the complete assignment names (with spaces at the end) from a partial
    assignment name from a particular class. This requires running `gkeep query assignments`.
    """
    data = __gkeep_query('assignments')
    if class_name not in data: data = __gkeep_query('assignments', 0.5)
    # Each assignment name starts with P for published or U for unpublished
    return __complete_word(assignment, [x.split(maxsplit=1)[1] for x in data.get(class_name, ())])

def __complete_student(class_name, student):
    """
    Return a list of the complete student names (with spaces at the end) from a partial student name
    from a particular class. This requires running `gkeep query students`.
    """
    data = __gkeep_query('students')
    if class_name not in data: data = __gkeep_query('students', 0.5)
    return __complete_word(student, data.get(class_name, ()))

def __add(words):
    """Completion for the command line gkeep add <new class name> <csv file>"""
    if len(words) == 2: return __complete_csv(words[1])
    return [] # either a new class name (which could be anything) or an unknown argument

def __modify(words):
    """Completion for the command line gkeep modify <class name> <csv file>"""
    if len(words) == 1: return __complete_class_name(words[0])
    if len(words) == 2: return __complete_csv(words[1])
    return [] # an unknown argument

def __upload(words):
    """Completion for the command line gkeep upload <class name> <directory>"""
    if len(words) == 1: return __complete_class_name(words[0])
    if len(words) == 2: return __complete_directory(words[1])
    return [] # an unknown argument

def __update(words):
    """
    Completion for the command line
    gkeep update <class name> <directory> { base_code | email | tests | all }
    """
    if len(words) == 1: return __complete_class_name(words[0])
    if len(words) == 2: return __complete_directory(words[1])
    if len(words) == 3: return __complete_word(words[2], ('base_code', 'email', 'tests', 'all'))
    return [] # an unknown argument

def __publish(words):
    """Completion for the command line gkeep publish <class name> <assignment>"""
    if len(words) == 1: return __complete_class_name(words[0])
    if len(words) == 2: return __complete_assignment(words[0], words[1])
    return [] # an unknown argument

def __delete(words):
    """Completion for the command line gkeep delete <class name> <assignment>"""
    if len(words) == 1: return __complete_class_name(words[0])
    if len(words) == 2: return __complete_assignment(words[0], words[1])
    return [] # an unknown argument

def __fetch(words):
    """Completion for the command line gkeep fetch <class name> <assignment> <directory>"""
    if len(words) == 1: return __complete_class_name(words[0])
    if len(words) == 2: return __complete_assignment(words[0], words[1])
    if len(words) == 3: return __complete_directory(words[2], False)
    return [] # an unknown argument

def __query(words):
    """
    Completion for the command line
    gkeep query { classes | assignments | recent | students } <#>
    """
    if len(words) == 1:
        return __complete_word(words[0], ('classes', 'assignments', 'recent', 'students'))
    return [] # either a number of days for recent or an unknown argument

def __trigger(words):
    """Completion for the command line gkeep trigger <class name> <assignment> <student> ..."""
    if len(words) == 1: return __complete_class_name(words[0])
    if len(words) == 2: return __complete_assignment(words[0], words[1])
    # List possible students but remove any already listed
    possible_students = __complete_student(words[0], words[-1])
    already_listed = set(words[2:-1])
    return [student for student in possible_students if student[:-1] not in already_listed]

def __status(words):
    """Completion for the command line gkeep status <class name> { open | closed }"""
    if len(words) == 1: return __complete_class_name(words[0])
    if len(words) == 2: return __complete_word(words[1], ('open', 'closed'))
    return [] # an unknown argument

__subcommands = {
    'add': __add,
    'modify': __modify,
    'upload': __upload,
    'update': __update,
    'publish': __publish,
    'delete': __delete,
    'fetch': __fetch,
    'query': __query,
    'trigger': __trigger,
    'config': None, # no arguments
    'status': __status,
    'add_faculty': None, # all arguments can be anything so just don't have a completion function
}

def command_completion(line):
    """
    Complete a gkeep command line. The user is currently typing at the end of the given line.
    Returns a list of the possiblities to fill in the final word with.
    """
    try:
        words = shlex.split(line)
    except ValueError:
        # We may have an unmatched quote in the final token
        # Try to add either a ' or " at the end
        try:
            words = shlex.split(line + '"')
        except ValueError:
            words = shlex.split(line + "'")

    # If we ended with an un-quoted space add a new blank word to the end
    if line[-1].isspace() and words[-1][-1] != line[-1]: words.append('')

    # Shouldn't happen, just the gkeep command...
    if len(words) == 1: return []

    # Name of one of the sub-commands
    if len(words) == 2: 
        return [key+' ' for key in __subcommands.keys() if key.startswith(words[1])]

    # Use the sub-command to do the remainder of the command line arguments
    __subcommand_completion = __subcommands.get(words[1])
    if __subcommand_completion is None: return []
    return __subcommand_completion(words[2:])

def __make_executable():
    """
    Makes sure this script is executable. Returns the command to be run by bash to register the
    completion for the current shell.
    """
    import stat
    script = os.path.abspath(__file__)
    
    # Make sure the script is executable
    try:
        mode = os.stat(script).st_mode
        os.chmod(script, mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
    except PermissionError: pass # let's hope it is already executable
    
    # Return the command to actually be run by bash
    return "complete -o nospace -C \'"+script+"\' gkeep"

def install():
    """
    Installs the code completion handler with bash-completions. The bash-completions package is
    required to be installed for this to be useful. If bash-completions is not installed then this
    will have no effect (well, it will write some very small files). If you want to use this
    without bash-completions run install_rc(), however this is the recommended way to have it
    installed.

    This follows the guidelines from https://github.com/scop/bash-completion/blob/master/README.md.

    If this is run as root, then the completion hook is installed into the global file. Otherwise it
    is instead in the user-specific directory.
    """
    cmd = __make_executable()
    
    # First get the bash-completions directory
    if os.geteuid() == 0:
        # We are root, install in a global directory
        info = subprocess.run(['pkg-config', '--variable=completionsdir', 'bash-completion'],
                              capture_output=True, text=True)
        if info.returncode == 0:
            directory = info.stdout.strip()
        else:
            directory = '/etc/bash_completion.d' # default location
            
    else:
        # Not root, install within user's profile
        if 'BASH_COMPLETION_USER_DIR' in os.environ:
            directory = os.environ['BASH_COMPLETION_USER_DIR']
        elif 'XDG_DATA_HOME' in os.environ:
            directory = os.path.join(os.environ['XDG_DATA_HOME'], 'bash-completion')
        elif os.path.isdir(os.path.expanduser('~/.local/share')):
            directory = os.path.expanduser('~/.local/share/bash-completion')
        else:
            # No viable directory to put it in, final choice is ~/.bash_completion
            with open(os.path.expanduser('~/.bash_completion'), 'a') as f:
                f.write('\n# Added by gkeep\n'+cmd+'\n')
            return # don't add to a directory
        
        directory = os.path.join(directory, 'completions')

    # Now add a file to register the complete command for gkeep
    os.makedirs(directory, exist_ok=True)
    with open(os.path.join(directory, 'gkeep'), 'w') as f:
        f.write('# bash completion for gkeep\n\n')
        f.write(cmd+'\n')

def install_rc():
    """
    Installs this completer in either ~/.bashrc or /etc/bashrc depending on if the current user it
    a regular user or root. The install() method is preferred to this method.
    """
    cmd = __make_executable()
    file = '/etc/bashrc' if os.geteuid() == 0 else os.path.expanduser('~/.bashrc')
    with open(file, 'a') as f:
        f.write('\n# Added by gkeep\n'+cmd+'\n')

def __main():
    if 'COMP_LINE' not in os.environ or 'COMP_POINT' not in os.environ:
        cmd = __make_executable()
        print("Program needs to be run from the programatic completion of bash", file=sys.stderr)
        print("To register it do `"+cmd+"` in bash", file=sys.stderr)
        sys.exit(1)

    # Get the command line up to the carat point
    point = int(os.environ['COMP_POINT'])
    line = os.environ['COMP_LINE'][:point]

    # Other environmental variables that could be used:
    #type_ = chr(int(os.environ['COMP_TYPE'])) # one of \t, ?, ! @, %
    #key = chr(int(os.environ['COMP_KEY']))

    # Get the completions
    completions = command_completion(line)

    # Print each option on its own line
    print('\n'.join(completions))

if __name__ == "__main__": __main()