MegEngine/imperative/python/megengine/tools/mge

#! /usr/bin/env python3
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
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:])