activate_pyzshcomplete

#!/usr/bin/python3

import os
from os.path import join, dirname, expanduser
from os import remove, geteuid, walk
from subprocess import run, Popen, PIPE
from sys import stderr, exit
from glob import glob


def main():
    try:
        avoid_running_as_super_user_implicitly()
        completion_script_dir = find_completion_script_dir()
        copy_completion_scripts(completion_script_dir)
        remove_completion_system_cache_files()
        print_success_message()
    except Exception as e:
        stderr.write('An error occured during installation: {}\n'.
            format(str(e)))
        exit(-1)


def avoid_running_as_super_user_implicitly():
    if geteuid() == 0:
        if input('You seem to be running as a super user. If you have a user '
                'account for which you would like to activate pyzshcomplete, '
                'you should run this script from that user. Root priviledges '
                'will be asked for when needed. \nContinue? [y/N] ').lower() \
                    != 'y':
            print('Aborting...')
            exit(0)
        print('Continuing...')


def find_completion_script_dir():
    print('Locating zsh completion script directory...')

    result = Popen(['echo $FPATH'], shell=True, executable='zsh', stdout=PIPE)
    fpath = result.stdout.read().decode('ascii').strip()

    for path in fpath.split(':'):
        for root, _, files in walk(path):
            if '_python' in files:
                print(f'Installing in {root}')
                return root

    raise RuntimeError('Can\'t find zsh completion folder. If you see this '
                        'error, please report it on the pyzshcomplete issue '
                        'tracker.\n')


def copy_completion_scripts(target_dir):
    print('Copying completion scripts... ')
    scripts = get_completion_scripts()
    perform_copy(scripts, target_dir, not os.access(target_dir, os.W_OK))


def remove_completion_system_cache_files():
    print('Removing completion system cache files... ')
    zsh_completion_cache_files = glob(expanduser('~/.zcompdump*'))
    for cache_file in zsh_completion_cache_files:
        try:
            remove(cache_file)
        except:
            stderr.write('WARNING: Failed removing zsh cache file: {}. You can '
                         'try to remove it manually.'.format(cache_file))


def print_success_message():
    print('pyzshcomplete was activated successfully!')
    print('Restart zsh for changes to take effect')


def get_completion_scripts():
    try:
        import pyzshcomplete
    except ImportError as e:
        raise ImportError('{}: Make sure the module is installed for the user '
                          'running the activation script'.format(str(e)))

    pyzshcomplete_root_dir = dirname(pyzshcomplete.__file__)
    zsh_scripts_dir = join(pyzshcomplete_root_dir, 'zsh_scripts')
    return ' '.join(glob(join(zsh_scripts_dir, '*')))


def perform_copy(source, target, as_root = False):
    commandline = f'{"sudo " if as_root else "" }cp {source} {target}'
    if run(commandline.split()).returncode != 0:
        raise RuntimeError('Failed copying zsh completion scripts. Make sure '
                           'to provide root access, as these files are copied '
                           'to a global zsh dir owned by root\n')


if __name__ == '__main__':
    main()