diff --git a/imperative/python/megengine/tools/mge b/imperative/python/megengine/tools/mge new file mode 100755 index 0000000000000000000000000000000000000000..e9db881e4620aa423f15b72e204f4c930b01e7c6 --- /dev/null +++ b/imperative/python/megengine/tools/mge @@ -0,0 +1,121 @@ +#! /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:]) diff --git a/imperative/python/setup.py b/imperative/python/setup.py index 99a9e2739a632e5bda3f0bc9da16350ba5643603..21c5365defce0d80059976c45e82a92b2437061a 100644 --- a/imperative/python/setup.py +++ b/imperative/python/setup.py @@ -106,6 +106,7 @@ setup_kwargs = dict( 'ci': requires_test, }, cmdclass={'build_ext': build_ext}, + scripts = ['./megengine/tools/mge'], )