feat: add cli autocomplete

jina/main/__init__.py

@@ -32,8 +32,22 @@ def _get_run_args(print_args: bool = True):
         exit()
 
 
+def _quick_ac_lookup():
+    from .autocomplete import ac_table
+    if sys.argv[1] == 'commands':
+        for k in ac_table['commands']:
+            print(k)
+        exit()
+    elif sys.argv[1] == 'completions':
+        if sys.argv[2] in ac_table['completions']:
+            for k in ac_table['completions'][sys.argv[2]]:
+                print(k)
+        exit()
+
+
 def main():
     """The main entrypoint of the CLI """
+    _quick_ac_lookup()
     from . import api
     args = _get_run_args()
     getattr(api, args.cli.replace('-', '_'))(args)

jina/main/api.py

@@ -1,6 +1,7 @@
 __copyright__ = "Copyright (c) 2020 Jina AI Limited. All rights reserved."
 __license__ = "Apache-2.0"
 
+
 def pod(args):
     """Start a Pod"""
     from ..peapods import Pod

jina/main/autocomplete.py

@@ -0,0 +1,39 @@
+def _update_autocomplete():
+    from jina.main.parser import get_main_parser, set_pea_parser, \
+        set_hw_parser, set_flow_parser, set_pod_parser, \
+        set_check_parser, set_gateway_parser, set_ping_parser, set_client_cli_parser, set_logger_parser
+
+    def _gaa(parser):
+        _compl = []
+        for v in parser._actions:
+            if v.option_strings:
+                _compl.extend(v.option_strings)
+            elif v.choices:
+                _compl.extend(v.choices)
+        # filer out single dash, as they serve as abbrev
+        _compl = [k for k in _compl if (not k.startswith('-') or k.startswith('--'))]
+        return _compl
+
+    compl = {
+        'commands': _gaa(get_main_parser()),
+        'completions': {
+            'pea': _gaa(set_pea_parser()),
+            'hello-world': _gaa(set_hw_parser()),
+            'flow': _gaa(set_flow_parser()),
+            'pod': _gaa(set_pod_parser()),
+            'check': _gaa(set_check_parser()),
+            'gateway': _gaa(set_gateway_parser()),
+            'ping': _gaa(set_ping_parser()),
+            'client': _gaa(set_client_cli_parser()),
+            'log': _gaa(set_logger_parser())
+        }
+    }
+
+    with open(__file__, 'a') as fp:
+        fp.write(f'\nac_table = {compl}\n')
+
+
+if __name__ == '__main__':
+    _update_autocomplete()
+
+ac_table = {'commands': ['--help', '--version', '--version-full', 'hello-world', 'pod', 'flow', 'gateway', 'ping', 'check', 'pea', 'log', 'client'], 'completions': {'pea': ['--help', '--version', '--version-full', '--name', '--identity', '--yaml-path', '--py-modules', '--image', '--entrypoint', '--pull-latest', '--volumes', '--port-in', '--port-out', '--host-in', '--host-out', '--socket-in', '--socket-out', '--port-ctrl', '--ctrl-with-ipc', '--timeout', '--timeout-ctrl', '--timeout-ready', '--dump-interval', '--exit-no-dump', '--read-only', '--separated-workspace', '--replica-id', '--check-version', '--array-in-pb', '--num-part', '--memory-hwm', '--runtime', '--max-idle-time', '--log-sse', '--log-remote', '--log-profile', '--override-exec-log', '--host', '--port-grpc', '--max-message-size', '--proxy'], 'hello-world': ['--help', '--version', '--version-full', '--workdir', '--logserver', '--shards', '--replicas', '--index-yaml-path', '--index-data-url', '--index-batch-size', '--query-yaml-path', '--query-data-url', '--query-batch-size', '--num-query', '--top-k'], 'flow': ['--help', '--version', '--version-full', '--yaml-path', '--logserver', '--logserver-config', '--optimize-level', '--output-type', '--output-path'], 'pod': ['--help', '--version', '--version-full', '--name', '--identity', '--yaml-path', '--py-modules', '--image', '--entrypoint', '--pull-latest', '--volumes', '--port-in', '--port-out', '--host-in', '--host-out', '--socket-in', '--socket-out', '--port-ctrl', '--ctrl-with-ipc', '--timeout', '--timeout-ctrl', '--timeout-ready', '--dump-interval', '--exit-no-dump', '--read-only', '--separated-workspace', '--replica-id', '--check-version', '--array-in-pb', '--num-part', '--memory-hwm', '--runtime', '--max-idle-time', '--log-sse', '--log-remote', '--log-profile', '--override-exec-log', '--host', '--port-grpc', '--max-message-size', '--proxy', '--replicas', '--polling', '--scheduling', '--reducing-yaml-path', '--shutdown-idle'], 'check': ['--help', '--version', '--version-full', '--summary-exec', '--summary-driver'], 'gateway': ['--help', '--version', '--version-full', '--name', '--identity', '--yaml-path', '--py-modules', '--image', '--entrypoint', '--pull-latest', '--volumes', '--port-in', '--port-out', '--host-in', '--host-out', '--socket-in', '--socket-out', '--port-ctrl', '--ctrl-with-ipc', '--timeout', '--timeout-ctrl', '--timeout-ready', '--dump-interval', '--exit-no-dump', '--read-only', '--separated-workspace', '--replica-id', '--check-version', '--array-in-pb', '--num-part', '--memory-hwm', '--runtime', '--max-idle-time', '--log-sse', '--log-remote', '--log-profile', '--override-exec-log', '--host', '--port-grpc', '--max-message-size', '--proxy', '--prefetch', '--prefetch-on-recv', '--allow-spawn'], 'ping': ['--help', '--version', '--version-full', '--timeout', '--retries', '--print-response'], 'client': ['--help', '--version', '--version-full', '--host', '--port-grpc', '--max-message-size', '--proxy', '--batch-size', '--mode', '--top-k', '--in-proto', '--callback-on-body', '--first-request-id', '--first-doc-id', '--random-doc-id', '--timeout-ready'], 'log': ['--help', '--version', '--version-full', '--groupby-regex', '--refresh-time']}}

jina/resources/completions/jina.bash

@@ -0,0 +1,17 @@
+
+_jina() {
+  COMPREPLY=()
+  local word="${COMP_WORDS[COMP_CWORD]}"
+
+  if [ "$COMP_CWORD" -eq 1 ]; then
+    COMPREPLY=( $(compgen -W "$(jina commands)" -- "$word") )
+  else
+    local words=("${COMP_WORDS[@]}")
+    unset words[0]
+    unset words[$COMP_CWORD]
+    local completions=$(jina completions "${words[@]}")
+    COMPREPLY=( $(compgen -W "$completions" -- "$word") )
+  fi
+}
+
+complete -F _jina jina

jina/resources/completions/jina.fish

@@ -0,0 +1,23 @@
+function __fish_jina_needs_command
+  set cmd (commandline -opc)
+  if [ (count $cmd) -eq 1 -a $cmd[1] = 'jina' ]
+    return 0
+  end
+  return 1
+end
+
+function __fish_jina_using_command
+  set cmd (commandline -opc)
+  if [ (count $cmd) -gt 1 ]
+    if [ $argv[1] = $cmd[2] ]
+      return 0
+    end
+  end
+  return 1
+end
+
+complete -f -c jina -n '__fish_jina_needs_command' -a '(jina commands)'
+for cmd in (jina commands)
+  complete -f -c jina -n "__fish_jina_using_command $cmd" -a \
+    "(jina completions (commandline -opc)[2..-1])"
+end

jina/resources/completions/jina.zsh

@@ -0,0 +1,18 @@
+if [[ ! -o interactive ]]; then
+    return
+fi
+
+compctl -K _jina jina
+
+_jina() {
+  local words completions
+  read -cA words
+
+  if [ "${#words}" -eq 2 ]; then
+    completions="$(jina commands)"
+  else
+    completions="$(jina completions ${words[2,-2]})"
+  fi
+
+  reply=(${(ps:\n:)completions})
+}

setup.py

@@ -1,7 +1,10 @@
 import sys
 from os import path
 
-from setuptools import setup, find_packages
+from setuptools import find_packages
+from setuptools import setup
+from setuptools.command.develop import develop
+from setuptools.command.install import install
 
 PY37 = 'py37'
 PY38 = 'py38'
@@ -67,6 +70,55 @@ def get_extra_requires(path, add_all=True):
         return {}
 
 
+def register_ac():
+    from pathlib import Path
+    import os
+    from pkg_resources import resource_filename
+
+    home = str(Path.home())
+    _check = [{'sh': os.path.join(home, '.zshrc'),
+               'ac': resource_filename('jina', '/'.join(('resources', 'completions', 'jina.zsh')))},
+              {'sh': os.path.join(home, '.bashrc'),
+               'ac': resource_filename('jina', '/'.join(('resources', 'completions', 'jina.bash')))},
+              {'sh': os.path.join(home, '.config', 'fish', 'config.fish'),
+               'ac': resource_filename('jina', '/'.join(('resources', 'completions', 'jina.fish')))}]
+
+    def add_ac(f):
+        if os.path.exists(f['sh']):
+            # zsh installed:
+            already_in = False
+            with open(f['sh']) as fp:
+                for v in fp:
+                    if f['ac'] in v:
+                        already_in = True
+                        break
+            if not already_in:
+                with open(f['sh'], 'a') as fp:
+                    fp.write('\nsource %s\n' % f['ac'])
+
+    try:
+        for k in _check:
+            add_ac(k)
+    except Exception:
+        pass
+
+
+class PostDevelopCommand(develop):
+    """Post-installation for development mode."""
+
+    def run(self):
+        develop.run(self)
+        register_ac()
+
+
+class PostInstallCommand(install):
+    """Post-installation for installation mode."""
+
+    def run(self):
+        install.run(self)
+        register_ac()
+
+
 setup(
     name=pkg_name,
     packages=find_packages(),
@@ -89,6 +141,10 @@ def get_extra_requires(path, add_all=True):
     entry_points={
         'console_scripts': ['jina=jina.main:main'],
     },
+    cmdclass={
+        'develop': PostDevelopCommand,
+        'install': PostInstallCommand,
+    },
     classifiers=(
         'Development Status :: 5 - Production/Stable',
         'Intended Audience :: Developers',