Add custom error and CustomHelpFormatter for cli (#888)

src/datachain/cli/parser/__init__.py

@@ -1,5 +1,4 @@
 import argparse
-from argparse import ArgumentParser
 from importlib.metadata import PackageNotFoundError, version
 
 import shtab
@@ -10,12 +9,16 @@
 from .studio import add_auth_parser
 from .utils import (
     FIND_COLUMNS,
+    CustomHelpFormatter,
     add_anon_arg,
     add_show_args,
     add_sources_arg,
     add_update_arg,
     find_columns_type,
 )
+from .utils import (
+    CustomArgumentParser as ArgumentParser,
+)
 
 
 def get_parser() -> ArgumentParser:  # noqa: PLR0915
@@ -28,10 +31,11 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     parser = ArgumentParser(
         description="DataChain: Wrangle unstructured AI data at scale.",
         prog="datachain",
+        formatter_class=CustomHelpFormatter,
     )
     parser.add_argument("-V", "--version", action="version", version=__version__)
 
-    parent_parser = ArgumentParser(add_help=False)
+    parent_parser = ArgumentParser(add_help=False, formatter_class=CustomHelpFormatter)
     parent_parser.add_argument(
         "-v", "--verbose", action="count", default=0, help="Be verbose"
     )
@@ -59,7 +63,10 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help=f"Use `{parser.prog} command --help` for command-specific help",
     )
     parse_cp = subp.add_parser(
-        "cp", parents=[parent_parser], description="Copy data files from the cloud."
+        "cp",
+        parents=[parent_parser],
+        description="Copy data files from the cloud.",
+        formatter_class=CustomHelpFormatter,
     )
     add_sources_arg(parse_cp).complete = shtab.DIR  # type: ignore[attr-defined]
     parse_cp.add_argument(
@@ -90,7 +97,10 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     add_update_arg(parse_cp)
 
     parse_clone = subp.add_parser(
-        "clone", parents=[parent_parser], description="Copy data files from the cloud."
+        "clone",
+        parents=[parent_parser],
+        description="Copy data files from the cloud.",
+        formatter_class=CustomHelpFormatter,
     )
     add_sources_arg(parse_clone).complete = shtab.DIR  # type: ignore[attr-defined]
     parse_clone.add_argument(
@@ -134,6 +144,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         aliases=["ds"],
         parents=[parent_parser],
         description="Commands for managing datasets.",
+        formatter_class=CustomHelpFormatter,
     )
     add_anon_arg(datasets_parser)
     datasets_subparser = datasets_parser.add_subparsers(
@@ -145,6 +156,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "pull",
         parents=[parent_parser],
         description="Pull specific dataset version from Studio.",
+        formatter_class=CustomHelpFormatter,
     )
     parse_pull.add_argument(
         "dataset",
@@ -188,7 +200,10 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_edit_dataset = datasets_subparser.add_parser(
-        "edit", parents=[parent_parser], description="Edit dataset metadata."
+        "edit",
+        parents=[parent_parser],
+        description="Edit dataset metadata.",
+        formatter_class=CustomHelpFormatter,
     )
     parse_edit_dataset.add_argument("name", type=str, help="Dataset name")
     parse_edit_dataset.add_argument(
@@ -234,7 +249,10 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     datasets_ls_parser = datasets_subparser.add_parser(
-        "ls", parents=[parent_parser], description="List datasets."
+        "ls",
+        parents=[parent_parser],
+        description="List datasets.",
+        formatter_class=CustomHelpFormatter,
     )
     datasets_ls_parser.add_argument(
         "--studio",
@@ -264,7 +282,11 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     rm_dataset_parser = datasets_subparser.add_parser(
-        "rm", parents=[parent_parser], description="Remove dataset.", aliases=["remove"]
+        "rm",
+        parents=[parent_parser],
+        description="Remove dataset.",
+        aliases=["remove"],
+        formatter_class=CustomHelpFormatter,
     )
     rm_dataset_parser.add_argument("name", type=str, help="Dataset name")
     rm_dataset_parser.add_argument(
@@ -308,7 +330,10 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_ls = subp.add_parser(
-        "ls", parents=[parent_parser], description="List storage contents."
+        "ls",
+        parents=[parent_parser],
+        description="List storage contents.",
+        formatter_class=CustomHelpFormatter,
     )
     add_anon_arg(parse_ls)
     add_update_arg(parse_ls)
@@ -348,7 +373,10 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_du = subp.add_parser(
-        "du", parents=[parent_parser], description="Display space usage."
+        "du",
+        parents=[parent_parser],
+        description="Display space usage.",
+        formatter_class=CustomHelpFormatter,
     )
     add_sources_arg(parse_du)
     add_anon_arg(parse_du)
@@ -380,7 +408,10 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_find = subp.add_parser(
-        "find", parents=[parent_parser], description="Search in a directory hierarchy."
+        "find",
+        parents=[parent_parser],
+        description="Search in a directory hierarchy.",
+        formatter_class=CustomHelpFormatter,
     )
     add_anon_arg(parse_find)
     add_update_arg(parse_find)
@@ -435,7 +466,10 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_index = subp.add_parser(
-        "index", parents=[parent_parser], description="Index storage location."
+        "index",
+        parents=[parent_parser],
+        description="Index storage location.",
+        formatter_class=CustomHelpFormatter,
     )
     add_anon_arg(parse_index)
     add_update_arg(parse_index)
@@ -445,6 +479,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "show",
         parents=[parent_parser],
         description="Create a new dataset with a query script.",
+        formatter_class=CustomHelpFormatter,
     )
     show_parser.add_argument("name", type=str, help="Dataset name")
     show_parser.add_argument(
@@ -461,6 +496,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "query",
         parents=[parent_parser],
         description="Create a new dataset with a query script.",
+        formatter_class=CustomHelpFormatter,
     )
     add_anon_arg(query_parser)
     query_parser.add_argument(
@@ -491,11 +527,15 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "clear-cache",
         parents=[parent_parser],
         description="Clear the local file cache.",
+        formatter_class=CustomHelpFormatter,
     )
     add_anon_arg(parse_clear_cache)
 
     parse_gc = subp.add_parser(
-        "gc", parents=[parent_parser], description="Garbage collect temporary tables."
+        "gc",
+        parents=[parent_parser],
+        description="Garbage collect temporary tables.",
+        formatter_class=CustomHelpFormatter,
     )
     add_anon_arg(parse_gc)
 
@@ -510,6 +550,7 @@ def add_completion_parser(subparsers, parents):
         "completion",
         parents=parents,
         description="Output shell completion script.",
+        formatter_class=CustomHelpFormatter,
     )
     parser.add_argument(
         "-s",

src/datachain/cli/parser/job.py

@@ -1,8 +1,15 @@
+from datachain.cli.parser.utils import CustomHelpFormatter
+
+
 def add_jobs_parser(subparsers, parent_parser) -> None:
     jobs_help = "Manage jobs in Studio"
     jobs_description = "Commands to manage job execution in Studio."
     jobs_parser = subparsers.add_parser(
-        "job", parents=[parent_parser], description=jobs_description, help=jobs_help
+        "job",
+        parents=[parent_parser],
+        description=jobs_description,
+        help=jobs_help,
+        formatter_class=CustomHelpFormatter,
     )
     jobs_subparser = jobs_parser.add_subparsers(
         dest="cmd",
@@ -17,10 +24,11 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         parents=[parent_parser],
         description=studio_run_description,
         help=studio_run_help,
+        formatter_class=CustomHelpFormatter,
     )
 
     studio_run_parser.add_argument(
-        "query_file",
+        "file",
         action="store",
         help="Query file to run",
     )
@@ -78,10 +86,11 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         parents=[parent_parser],
         description=studio_cancel_description,
         help=studio_cancel_help,
+        formatter_class=CustomHelpFormatter,
     )
 
     studio_cancel_parser.add_argument(
-        "job_id",
+        "id",
         action="store",
         help="Job ID to cancel",
     )
@@ -100,10 +109,11 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         parents=[parent_parser],
         description=studio_log_description,
         help=studio_log_help,
+        formatter_class=CustomHelpFormatter,
     )
 
     studio_log_parser.add_argument(
-        "job_id",
+        "id",
         action="store",
         help="Job ID to show logs for",
     )

src/datachain/cli/parser/studio.py

@@ -1,3 +1,6 @@
+from datachain.cli.parser.utils import CustomHelpFormatter
+
+
 def add_auth_parser(subparsers, parent_parser) -> None:
     from dvc_studio_client.auth import AVAILABLE_SCOPES
 
@@ -9,6 +12,7 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         parents=[parent_parser],
         description=auth_description,
         help=auth_help,
+        formatter_class=CustomHelpFormatter,
     )
     auth_subparser = auth_parser.add_subparsers(
         dest="cmd",
@@ -27,6 +31,7 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         parents=[parent_parser],
         description=auth_login_description,
         help=auth_login_help,
+        formatter_class=CustomHelpFormatter,
     )
 
     login_parser.add_argument(
@@ -69,6 +74,7 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         parents=[parent_parser],
         description=auth_logout_description,
         help=auth_logout_help,
+        formatter_class=CustomHelpFormatter,
     )
 
     auth_team_help = "Set default team for Studio operations"
@@ -79,6 +85,7 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         parents=[parent_parser],
         description=auth_team_description,
         help=auth_team_help,
+        formatter_class=CustomHelpFormatter,
     )
     team_parser.add_argument(
         "team_name",
@@ -100,4 +107,5 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         parents=[parent_parser],
         description=auth_token_description,
         help=auth_token_help,
+        formatter_class=CustomHelpFormatter,
     )

src/datachain/cli/parser/utils.py

@@ -1,11 +1,30 @@
-from argparse import Action, ArgumentParser, ArgumentTypeError
+from argparse import Action, ArgumentParser, ArgumentTypeError, HelpFormatter
 from typing import Union
 
 from datachain.cli.utils import CommaSeparatedArgs
 
 FIND_COLUMNS = ["du", "name", "path", "size", "type"]
 
 
+class CustomHelpFormatter(HelpFormatter):
+    def add_arguments(self, actions):
+        # Sort arguments to move --help and --version to the end
+        normal_actions = [
+            a for a in actions if a.dest not in ("help", "verbose", "quiet")
+        ]
+        special_actions = [a for a in actions if a.dest in ("help", "verbose", "quiet")]
+        super().add_arguments(normal_actions + special_actions)
+
+
+class CustomArgumentParser(ArgumentParser):
+    def error(self, message):
+        internal_commands = ["internal-run-udf", "internal-run-udf-worker"]
+
+        hidden_portion = "".join(f"'{cmd}', " for cmd in internal_commands)
+        message = message.replace(hidden_portion, "")
+        super().error(message)
+
+
 def find_columns_type(
     columns_str: str,
     default_colums_str: str = "path",

src/datachain/studio.py

@@ -29,7 +29,7 @@ def process_jobs_args(args: "Namespace"):
 
     if args.cmd == "run":
         return create_job(
-            args.query_file,
+            args.file,
             args.team,
             args.env_file,
             args.env,
@@ -41,9 +41,9 @@ def process_jobs_args(args: "Namespace"):
         )
 
     if args.cmd == "cancel":
-        return cancel_job(args.job_id, args.team)
+        return cancel_job(args.id, args.team)
     if args.cmd == "logs":
-        return show_job_logs(args.job_id, args.team)
+        return show_job_logs(args.id, args.team)
     raise DataChainError(f"Unknown command '{args.cmd}'.")
 
 

tests/unit/test_cli_parsing.py

@@ -1,12 +1,13 @@
 import logging
-from argparse import ArgumentParser, ArgumentTypeError
+from argparse import ArgumentTypeError
 
 import pytest
 
 from datachain.cli import (
     get_logging_level,
     get_parser,
 )
+from datachain.cli.parser.utils import CustomArgumentParser as ArgumentParser
 from datachain.cli.parser.utils import find_columns_type
 from datachain.cli.utils import CommaSeparatedArgs, KeyValueArgs