From d849bd3b66c2af7802a3be62b116f2f9ff8099ce Mon Sep 17 00:00:00 2001
From: Kelly Brazil <kellyjonbrazil@gmail.com>
Date: Sun, 5 Jun 2022 16:23:55 -0700
Subject: [PATCH] add bash and zsh completions to cli

---
 jc/cli.py               | 25 +++++++------
 jc/cli_data.py          | 18 ++++++++++
 jc/shell_completions.py | 79 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 109 insertions(+), 13 deletions(-)
 create mode 100644 jc/cli_data.py
 create mode 100644 jc/shell_completions.py

diff --git a/jc/cli.py b/jc/cli.py
index 927b25fe..c74c73be 100644
--- a/jc/cli.py
+++ b/jc/cli.py
@@ -14,6 +14,8 @@ from .lib import (__version__, parser_info, all_parser_info, parsers,
                   _get_parser, _parser_is_streaming, standard_parser_mod_list,
                   plugin_parser_mod_list, streaming_parser_mod_list)
 from . import utils
+from .cli_data import long_options_map
+from .shell_completions import bash_completion, zsh_completion
 from . import tracebackplus
 from .exceptions import LibraryNotInstalled, ParseError
 
@@ -84,19 +86,6 @@ if PYGMENTS_INSTALLED:
             'white': 'ansiwhite',
         }
 
-long_options_map: Dict[str, List[str]] = {
-    '--about': ['a', 'about jc'],
-    '--force-color': ['C', 'force color output even when using pipes (overrides -m)'],
-    '--debug': ['d', 'debug (double for verbose debug)'],
-    '--help': ['h', 'help (--help --parser_name for parser documentation)'],
-    '--monochrome': ['m', 'monochrome output'],
-    '--pretty': ['p', 'pretty print output'],
-    '--quiet': ['q', 'suppress warnings (double to ignore streaming errors)'],
-    '--raw': ['r', 'raw output'],
-    '--unbuffer': ['u', 'unbuffer output'],
-    '--version': ['v', 'version info'],
-    '--yaml-out': ['y', 'YAML output']
-}
 
 def set_env_colors(env_colors=None):
     """
@@ -530,6 +519,8 @@ def main():
     unbuffer = 'u' in options
     version_info = 'v' in options
     yaml_out = 'y' in options
+    bash_comp = 'B' in options
+    zsh_comp = 'Z' in options
 
     if verbose_debug:
         tracebackplus.enable(context=11)
@@ -554,6 +545,14 @@ def main():
         utils._safe_print(versiontext())
         sys.exit(0)
 
+    if bash_comp:
+        utils._safe_print(bash_completion())
+        sys.exit(0)
+
+    if zsh_comp:
+        utils._safe_print(zsh_completion())
+        sys.exit(0)
+
     # if magic syntax used, try to run the command and error if it's not found, etc.
     magic_stdout, magic_stderr, magic_exit_code = None, None, 0
     if run_command:
diff --git a/jc/cli_data.py b/jc/cli_data.py
new file mode 100644
index 00000000..f9f6ed2f
--- /dev/null
+++ b/jc/cli_data.py
@@ -0,0 +1,18 @@
+"""jc - JSON Convert cli_data module"""
+from typing import List, Dict
+
+long_options_map: Dict[str, List[str]] = {
+    '--about': ['a', 'about jc'],
+    '--force-color': ['C', 'force color output even when using pipes (overrides -m)'],
+    '--debug': ['d', 'debug (double for verbose debug)'],
+    '--help': ['h', 'help (--help --parser_name for parser documentation)'],
+    '--monochrome': ['m', 'monochrome output'],
+    '--pretty': ['p', 'pretty print output'],
+    '--quiet': ['q', 'suppress warnings (double to ignore streaming errors)'],
+    '--raw': ['r', 'raw output'],
+    '--unbuffer': ['u', 'unbuffer output'],
+    '--version': ['v', 'version info'],
+    '--yaml-out': ['y', 'YAML output'],
+    '--bash-comp': ['B', 'gen Bash completion: jc -B > /etc/bash_completion.d/jc'],
+    '--zsh-comp': ['Z', 'gen Zsh completion: jc -Z > "${fpath[1]}/_jc"']
+}
diff --git a/jc/shell_completions.py b/jc/shell_completions.py
new file mode 100644
index 00000000..d5567959
--- /dev/null
+++ b/jc/shell_completions.py
@@ -0,0 +1,79 @@
+"""jc - JSON Convert shell_completions module"""
+
+from string import Template
+from .cli_data import long_options_map
+from .lib import all_parser_info
+
+# $(jc -a | jq -r '.parsers[] | .argument, .magic_commands[]?')
+bash_template = Template('''\
+complete -W "${bash_arguments}${bash_options}${bash_commands}" jc
+''')
+
+zsh_template = Template('''\
+#compdef jc
+
+_jc() {
+  # autogenerate completions based on jc --help output
+  _arguments --
+
+  # add commands supported by magic syntax
+  local -a commands
+  commands=(
+    # e.g. 'arp:run arp with magic syntax.'
+    ${zsh_commands}
+  )
+
+  _describe -t commands 'commands' commands
+  return 0
+}
+
+_jc
+''')
+
+def get_commands():
+    command_list = []
+    for cmd in all_parser_info():
+        if 'magic_commands' in cmd:
+            command_list.extend(cmd['magic_commands'])
+
+    return list(set([i.split()[0] for i in command_list]))
+
+
+def get_options():
+    options_list = []
+    for opt in long_options_map:
+        options_list.append(opt)
+        options_list.append('-' + long_options_map[opt][0])
+
+    return options_list
+
+
+def get_arguments():
+    arg_list = []
+    for cmd in all_parser_info():
+        if 'argument' in cmd:
+            arg_list.append(cmd['argument'])
+
+    return arg_list
+
+
+def gen_zsh_command_descriptions(command_list):
+    zsh_commands = []
+    for cmd in command_list:
+        zsh_commands.append(f"""'{cmd}:run "{cmd}" command with magic syntax.'""")
+
+    return zsh_commands
+
+
+def bash_completion():
+    args = '\n'.join(get_arguments())
+    options = '\n' + '\n'.join(get_options())
+    commands = '\n' + '\n'.join(get_commands())
+    return bash_template.substitute(bash_arguments=args,
+                                    bash_options=options,
+                                    bash_commands=commands)
+
+
+def zsh_completion():
+    commands = '\n    '.join(gen_zsh_command_descriptions(get_commands()))
+    return zsh_template.substitute(zsh_commands=commands)