MegEngine/imperative/python/megengine/tools/mge

#! /usr/bin/env python3
import argparse
import ntpath
import os
import pathlib
import platform
import sys
from glob import glob

import megengine


def add_completion(tools_path, write_path) :
    script_files = [os.path.basename(p) for p in tools_path.glob("*.py")]
    param_script = []
    for script_name in script_files:
        temp = ntpath.splitext(script_name)[0]
        if temp=="__init__" or temp=="mge":
            continue
        param_script.append(temp)

    completion_template_head = '''
    _mge(){
        local cur prev script_name
        COMPREPLY=()
        cur="${COMP_WORDS[COMP_CWORD]}"
        prev="${COMP_WORDS[COMP_CWORD-1]}"
        case "${prev}" in
            mge)
            words="compare_binary_iodump draw_graph load_network_and_run network_visualize profile_analyze"
            COMPREPLY=( $(compgen -W "${words}" -- ${cur} ) )
            return
            ;;
    '''
    completion_template_body = ""
    completion_template_tail2=""
    for param in param_script:
        words_command =  'grep -Eso \'"-[[:alnum:]-]*[_[:alnum:]]*"\' {}/{}.py | xargs'.format(str(tools_path),param)
        words = os.popen(words_command).read().strip()

        completion_template_body+= '''
        {})
            words="{}"
            COMPREPLY=($(compgen -W "$words" -- $cur))
            return
            ;;
        '''.format(param ,words)
        completion_template_tail2+= '''
            {})
                words="{}"
                COMPREPLY=($(compgen -W "$words" -- $cur))
                return
                ;;
        '''.format(param ,words)
    completion_template_tail1='''
        esac
        case $cur in
            -*)
            script_name="${COMP_WORDS[1]}"
            case $script_name in
    '''

    completion_template_tail3='''
            esac
            ;;
        esac
        return
    }
    complete -o bashdefault -F _mge mge
    '''
    completion_template =  completion_template_head + completion_template_body + completion_template_tail1 + completion_template_tail2 + completion_template_tail3

    wp  = pathlib.Path(write_path)
    wp.parent.mkdir(parents=True, exist_ok = True)

    with open(write_path, 'w+') as f:
        f.write(completion_template)

def init(path) :
    engine_path = pathlib.Path(megengine.__path__[0])
    add_completion(engine_path/'tools',path)

def main(argv) :
    if len(argv) == 0:
        return
    script_file = argv[0]
    args = " ".join(argv[1:])
    call_command = "python3 -m megengine.tools.{} {}".format(script_file,args)
    os.system(call_command)

if __name__ == "__main__":
    usage = 'usage: mge [-h|--help|--h] [--init] [script_name --script_param xxx]'

    if len(sys.argv) <= 1 or sys.argv[1] == '-h' or sys.argv[1] =='--help' or sys.argv[1] =='--h':
        print(usage)
        sys.exit(0)
    if sys.argv[1] == '--init':
        sysstr = platform.system()
        if(sysstr == "Windows"):
            print("WARNING: windows doesn't support hinting")
        else:
            path = "{}/.local/share/bash-completion/completions/mge".format(os.environ.get("HOME"))
            init(path)
            shell = os.environ.get("SHELL")
            if shell.find('zsh') != -1:
                print("if you don't have zsh completion init, please excute command: 'autoload -Uz compinit && compinit'")
                print("Guess you are using zsh, please add 'source %s' to your ~/.zshrc" % path)
            elif shell.find('bash') != -1:
                print("Guess you are using bash, please relogin or do 'source %s'" % path)
            else:
                print("Current {} doesn't support hinting shell completion".format(shell))
            sys.exit(0)
    else:
        main(sys.argv[1:])