Refactor the CLI to be separate commands

.pre-commit-config.yaml

@@ -48,6 +48,7 @@ repos:
         - pytest
         - requests
         - tomli; python_version<"3.11"
+        - setuptools
 
   - repo: https://github.com/gitleaks/gitleaks
     rev: v8.21.2

command_line_assistant/__main__.py

@@ -1,46 +1,39 @@
-import logging
-import os
 import sys
 from pathlib import Path
 
-from command_line_assistant.cli import get_args
+from command_line_assistant.commands import history, query, record
 from command_line_assistant.config import (
+    CONFIG_DEFAULT_PATH,
     load_config_file,
 )
-from command_line_assistant.handlers import (
-    handle_history_write,
-    handle_query,
-    handle_script_session,
-)
 from command_line_assistant.logger import setup_logging
+from command_line_assistant.utils.cli import add_default_command, create_argument_parser
 
 
 def main() -> int:
-    parser, args = get_args()
-
-    config_file = Path(args.config).expanduser()
+    config_file = Path(CONFIG_DEFAULT_PATH)
     config = load_config_file(config_file)
 
-    setup_logging(config, args.verbose)
+    setup_logging(config, False)
+
+    parser, commands_parser = create_argument_parser()
 
-    enforce_script_session = config.output.enforce_script
-    output_file = config.output.file
+    # TODO: add autodetection of BaseCLICommand classes in the future so we can just drop
+    # new subcommand python modules into the directory and then loop and call `register_subcommand()`
+    # on each one.
+    query.register_subcommand(commands_parser, config)  # type: ignore
+    history.register_subcommand(commands_parser, config)  # type: ignore
+    record.register_subcommand(commands_parser, config)  # type: ignore
 
-    if enforce_script_session and (not args.record or not os.path.exists(output_file)):
-        parser.error(
-            f"Please call `{sys.argv[0]} --record` first to initialize script session or create the output file."
-        )
+    args = add_default_command(sys.argv)
+    args = parser.parse_args(args)
 
-    # NOTE: This needs more refinement, script session can't be combined with other arguments
-    if args.record:
-        handle_script_session(output_file)
-        return 0
-    if args.history_clear:
-        logging.info("Clearing history of conversation")
-        handle_history_write(config, [], "")
-    if args.query_string:
-        handle_query(args.query_string, config)
+    if not hasattr(args, "func"):
+        parser.print_help()
+        return 1
 
+    service = args.func(args)
+    service.run()
     return 0
 
 

command_line_assistant/commands/history.py

@@ -0,0 +1,43 @@
+import logging
+from argparse import Namespace
+
+from command_line_assistant.config import Config
+from command_line_assistant.history import handle_history_write
+from command_line_assistant.utils.cli import BaseCLICommand, SubParsersAction
+
+logger = logging.getLogger(__name__)
+
+
+class HistoryCommand(BaseCLICommand):
+    def __init__(self, clear: bool, config: Config) -> None:
+        self._clear = clear
+        self._config = config
+        super().__init__()
+
+    def run(self) -> None:
+        if self._clear:
+            logger.info("Clearing history of conversation")
+            handle_history_write(self._config, [], "")
+
+
+def register_subcommand(parser: SubParsersAction, config: Config):
+    """
+    Register this command to argparse so it's available for the datasets-cli
+
+    Args:
+        parser: Root parser to register command-specific arguments
+    """
+    history_parser = parser.add_parser(
+        "history",
+        help="Manage conversation history",
+    )
+    history_parser.add_argument(
+        "--clear", action="store_true", help="Clear the history."
+    )
+
+    # TODO(r0x0d): This is temporary as it will get removed
+    history_parser.set_defaults(func=lambda args: _command_factory(args, config))
+
+
+def _command_factory(args: Namespace, config: Config) -> HistoryCommand:
+    return HistoryCommand(args.clear, config)

command_line_assistant/commands/query.py

@@ -0,0 +1,39 @@
+from argparse import Namespace
+
+from command_line_assistant.config import Config
+from command_line_assistant.handlers import handle_query
+from command_line_assistant.utils.cli import BaseCLICommand, SubParsersAction
+
+
+class QueryCommand(BaseCLICommand):
+    def __init__(self, query_string: str, config: Config) -> None:
+        self._query = query_string
+        self._config = config
+        super().__init__()
+
+    def run(self) -> None:
+        handle_query(self._query, self._config)
+
+
+def register_subcommand(parser: SubParsersAction, config: Config) -> None:
+    """
+    Register this command to argparse so it's available for the datasets-cli
+
+    Args:
+        parser: Root parser to register command-specific arguments
+    """
+    query_parser = parser.add_parser(
+        "query",
+        help="",
+    )
+    # Positional argument, required only if no optional arguments are provided
+    query_parser.add_argument(
+        "query_string", nargs="?", help="Query string to be processed."
+    )
+
+    # TODO(r0x0d): This is temporary as it will get removed
+    query_parser.set_defaults(func=lambda args: _command_factory(args, config))
+
+
+def _command_factory(args: Namespace, config: Config) -> QueryCommand:
+    return QueryCommand(args.query_string, config)

command_line_assistant/commands/record.py

@@ -0,0 +1,49 @@
+import logging
+import os
+import sys
+
+from command_line_assistant.config import Config
+from command_line_assistant.handlers import handle_script_session
+from command_line_assistant.utils.cli import BaseCLICommand, SubParsersAction
+
+logger = logging.getLogger(__name__)
+
+# NOTE: This needs more refinement, script session can't be combined with other arguments
+
+
+class RecordCommand(BaseCLICommand):
+    def __init__(self, config: Config) -> None:
+        self._config = config
+        super().__init__()
+
+    def run(self) -> None:
+        enforce_script_session = self._config.output.enforce_script
+        output_file = self._config.output.file
+
+        if enforce_script_session and not os.path.exists(output_file):
+            logger.error(
+                "Please call `%s record` first to initialize script session or create the output file.",
+                sys.argv[0],
+            )
+
+        handle_script_session(output_file)
+
+
+def register_subcommand(parser: SubParsersAction, config: Config):
+    """
+    Register this command to argparse so it's available for the datasets-cli
+
+    Args:
+        parser: Root parser to register command-specific arguments
+    """
+    record_parser = parser.add_parser(
+        "record",
+        help="Start a recording session for script output.",
+    )
+
+    # TODO(r0x0d): This is temporary as it will get removed
+    record_parser.set_defaults(func=lambda args: _command_factory(config))
+
+
+def _command_factory(config: Config) -> RecordCommand:
+    return RecordCommand(config)

command_line_assistant/config.py

@@ -13,8 +13,9 @@
 except ImportError:
     import tomli as tomllib  # pyright: ignore[reportMissingImports]
 
-
-CONFIG_DEFAULT_PATH: Path = Path("~/.config/command-line-assistant/config.toml")
+CONFIG_DEFAULT_PATH: Path = Path(
+    "~/.config/command-line-assistant/config.toml"
+).expanduser()
 
 # tomllib does not support writting files, so we will create our own.
 CONFIG_TEMPLATE = """\

command_line_assistant/constants.py

@@ -0,0 +1 @@
+VERSION = "0.1.0"

command_line_assistant/utils/cli.py

@@ -0,0 +1,99 @@
+import select
+import sys
+from abc import ABC, abstractmethod
+from argparse import SUPPRESS, ArgumentParser, _SubParsersAction
+from typing import Optional
+
+from command_line_assistant.config import CONFIG_DEFAULT_PATH
+from command_line_assistant.constants import VERSION
+
+# Define the type here so pyright is happy with it.
+SubParsersAction = _SubParsersAction
+
+PARENT_ARGS: list[str] = ["--version", "-v", "-h", "--help"]
+ARGS_WITH_VALUES: list[str] = ["--clear"]
+
+
+class BaseCLICommand(ABC):
+    @abstractmethod
+    def run(self):
+        raise NotImplementedError("Not implemented in base class.")
+
+
+def add_default_command(argv):
+    """Add the default command when none is given"""
+    args = argv[1:]
+
+    # Early exit if we don't have any argv
+    if not args:
+        return args
+
+    subcommand = _subcommand_used(argv)
+    if subcommand is None:
+        args.insert(0, "query")
+
+    return args
+
+
+def _subcommand_used(args):
+    """Return what subcommand has been used by the user. Return None if no subcommand has been used."""
+    for index, argument in enumerate(args):
+        # If we have a exact match for any of the commands, return directly
+        if argument in ("query", "history"):
+            return argument
+
+        # It means that we hit a --version/--help
+        if argument in PARENT_ARGS:
+            return argument
+
+        # Otherwise, check if this is the second part of an arg that takes a value.
+        elif args[index - 1] in ARGS_WITH_VALUES:
+            continue
+
+    return None
+
+
+def create_argument_parser() -> tuple[ArgumentParser, SubParsersAction]:
+    """Create the argument parser for command line assistant."""
+    parser = ArgumentParser(
+        description="A script with multiple optional arguments and a required positional argument if no optional arguments are provided.",
+    )
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=VERSION,
+        default=SUPPRESS,
+        help="Show command line assistant version and exit.",
+    )
+    parser.add_argument(
+        "--config",
+        default=CONFIG_DEFAULT_PATH,
+        help="Path to the config file.",
+    )
+    commands_parser = parser.add_subparsers(
+        dest="command", help="command line assistant helpers"
+    )
+
+    return parser, commands_parser
+
+
+def read_stdin() -> Optional[str]:
+    """Parse the std input when a user give us.
+
+    For example, consider the following scenario:
+        >>> echo "how to run podman?" | c
+
+    Or a more complex one
+        >>> cat error-log | c "How to fix this?"
+
+    Returns:
+        In case we have a stdin, we parse and retrieve it. Otherwise, just
+        return None.
+    """
+    # Check if there's input available on stdin
+    if select.select([sys.stdin], [], [], 0.0)[0]:
+        # If there is input, read it
+        input_data = sys.stdin.read().strip()
+        return input_data
+    # If no input, return None or handle as you prefer
+    return None