kresctl: tab-completion: implement suggestions/completion for first argument

3 files changed, 75 insertions, 43 deletions

diff --git a/python/knot_resolver/client/command.py b/python/knot_resolver/client/command.py
index 76c0f1d0..c3055844 100644
--- a/python/knot_resolver/client/command.py
+++ b/python/knot_resolver/client/command.py
@@ -1,7 +1,7 @@
 import argparse
 from abc import ABC, abstractmethod  # pylint: disable=[no-name-in-module]
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Type, TypeVar
+from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar
 from urllib.parse import quote
 
 from knot_resolver.constants import API_SOCK_FILE, CONFIG_FILE
@@ -17,6 +17,33 @@ CompWords = Dict[str, Optional[str]]
 _registered_commands: List[Type["Command"]] = []
 
 
+def get_subparsers_words(subparser_actions: List[argparse.Action]) -> CompWords:
+    words: CompWords = {}
+    for action in subparser_actions:
+        if isinstance(action, argparse._SubParsersAction) and action.choices:  # noqa: SLF001
+            for choice, parser in action.choices.items():
+                words[choice] = parser.description
+        else:
+            for opt in action.option_strings:
+                words[opt] = action.help
+    return words
+
+
+def get_subparser_by_name(name: str, parser_actions: List[argparse.Action]) -> Optional[argparse.ArgumentParser]:
+    for action in parser_actions:
+        if isinstance(action, argparse._SubParsersAction):  # noqa: SLF001
+            if action.choices and name in action.choices:
+                return action.choices[name]
+    return None
+
+
+def get_subparser_command(subparser: argparse.ArgumentParser) -> "Command":
+    defaults: Dict[str, Any] = subparser._defaults  # noqa: SLF001
+    if "command" in defaults:
+        return defaults["command"]
+    raise ValueError(f"missing 'command' default for '{subparser.prog}' parser")
+
+
 def register_command(cls: T) -> T:
     _registered_commands.append(cls)
     return cls
diff --git a/python/knot_resolver/client/commands/cache.py b/python/knot_resolver/client/commands/cache.py
index 60417eec..d1116580 100644
--- a/python/knot_resolver/client/commands/cache.py
+++ b/python/knot_resolver/client/commands/cache.py
@@ -3,7 +3,7 @@ import sys
 from enum import Enum
 from typing import Any, Dict, List, Optional, Tuple, Type
 
-from knot_resolver.client.command import Command, CommandArgs, CompWords, register_command
+from knot_resolver.client.command import Command, CommandArgs, CompWords, get_subparsers_words, register_command
 from knot_resolver.datamodel.cache_schema import CacheClearRPCSchema
 from knot_resolver.utils.modeling.exceptions import AggregateDataValidationError, DataValidationError
 from knot_resolver.utils.modeling.parsing import DataFormat, parse_json
@@ -99,7 +99,7 @@ class CacheCommand(Command):
 
     @staticmethod
     def completion(args: List[str], parser: argparse.ArgumentParser) -> CompWords:
-        return {}
+        return get_subparsers_words(parser._actions)  # noqa: SLF001
 
     def run(self, args: CommandArgs) -> None:
         if not self.operation:
diff --git a/python/knot_resolver/client/commands/completion.py b/python/knot_resolver/client/commands/completion.py
index 05fdded8..ee46bc45 100644
--- a/python/knot_resolver/client/commands/completion.py
+++ b/python/knot_resolver/client/commands/completion.py
@@ -2,7 +2,15 @@ import argparse
 from enum import Enum
 from typing import List, Tuple, Type
 
-from knot_resolver.client.command import Command, CommandArgs, CompWords, register_command
+from knot_resolver.client.command import (
+    Command,
+    CommandArgs,
+    CompWords,
+    get_subparser_by_name,
+    get_subparser_command,
+    get_subparsers_words,
+    register_command,
+)
 
 
 class Shells(Enum):
@@ -25,7 +33,10 @@ class CompletionCommand(Command):
     def register_args_subparser(
         subparser: "argparse._SubParsersAction[argparse.ArgumentParser]",
     ) -> Tuple[argparse.ArgumentParser, "Type[Command]"]:
-        completion = subparser.add_parser("completion", help="commands auto-completion")
+        completion = subparser.add_parser(
+            "completion",
+            help="commands auto-completion",
+        )
         completion.add_argument(
             "--space",
             help="space after last word, returns all possible folowing options",
@@ -49,47 +60,41 @@ class CompletionCommand(Command):
 
     @staticmethod
     def completion(args: List[str], parser: argparse.ArgumentParser) -> CompWords:
-        words: CompWords = {}
-        # for action in parser._actions:
-        #     for opt in action.option_strings:
-        #         words[opt] = action.help
-        # return words
-        return words
+        return get_subparsers_words(parser._actions)  # noqa: SLF001
 
     def run(self, args: CommandArgs) -> None:
-        pass
-        # subparsers = args.parser._subparsers
-        # words: CompWords = {}
+        subparsers = args.parser._subparsers  # noqa: SLF001
+        words: CompWords = {}
 
-        # if subparsers:
-        #     words = parser_words(subparsers._actions)
+        if subparsers:
+            words = get_subparsers_words(subparsers._actions)  # noqa: SLF001
 
-        #     uargs = iter(self.comp_args)
-        #     for uarg in uargs:
-        #         subparser = subparser_by_name(uarg, subparsers._actions)  # pylint: disable=W0212
+            uargs = iter(self.comp_args)
+            for uarg in uargs:
+                subparser = get_subparser_by_name(uarg, subparsers._actions)  # noqa: SLF001
 
-        #         if subparser:
-        #             cmd: Command = subparser_command(subparser)
-        #             subparser_args = self.comp_args[self.comp_args.index(uarg) + 1 :]
-        #             if subparser_args:
-        #                 words = cmd.completion(subparser_args, subparser)
-        #             break
-        #         elif uarg in ["-s", "--socket"]:
-        #             # if arg is socket config, skip next arg
-        #             next(uargs)
-        #             continue
-        #         elif uarg in words:
-        #             # uarg is walid arg, continue
-        #             continue
-        #         else:
-        #             raise ValueError(f"unknown argument: {uarg}")
+                if subparser:
+                    cmd: Command = get_subparser_command(subparser)
+                    subparser_args = self.comp_args[self.comp_args.index(uarg) + 1 :]
+                    if subparser_args or self.space:
+                        words = cmd.completion(subparser_args, subparser)
+                    break
+                if uarg in ["-s", "--socket", "-c", "--config"]:
+                    # if arg is socket config, skip next arg
+                    next(uargs)
+                    continue
+                if uarg in words:
+                    # uarg is valid (complete) arg, continue
+                    continue
+                else:
+                    raise ValueError(f"unknown argument: {uarg}")
 
-        # # print completion words
-        # # based on required bash/fish shell format
-        # if self.shell == Shells.BASH:
-        #     print(" ".join(words))
-        # elif self.shell == Shells.FISH:
-        #     # TODO: FISH completion implementation
-        #     pass
-        # else:
-        #     raise ValueError(f"unexpected value of {Shells}: {self.shell}")
+        # print completion words
+        # based on required bash/fish shell format
+        if self.shell == Shells.BASH:
+            print(" ".join(words))
+        elif self.shell == Shells.FISH:
+            # TODO: FISH completion implementation
+            pass
+        else:
+            raise ValueError(f"unexpected value of {Shells}: {self.shell}")