feat(mge/tools): add bash_completion for user to hint script

BREAKING CHANGE:

fix(mge/tools): change help script for other shell

BREAKING CHANGE:

fix(mge/tools): fix double -h error

BREAKING CHANGE:

fix(mge/tools): add support for _ parameter

BREAKING CHANGE:

GitOrigin-RevId: 537c3a2aae3f13100f40543112a42e4f831a2e41

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:])

imperative/python/setup.py

@@ -106,6 +106,7 @@ setup_kwargs = dict(
         'ci': requires_test,
     },
     cmdclass={'build_ext': build_ext},
+    scripts = ['./megengine/tools/mge'],
 )