fix: simplify

Author: tobiasraabe

src/_pytask/click.py

@@ -33,7 +33,7 @@
     from collections.abc import Sequence
 
 
-__all__ = ["ColoredCommand", "ColoredGroup", "EnumChoice"]
+__all__ = ["ColoredCommand", "ColoredGroup", "EnumChoice", "get_command"]
 
 
 if importlib.metadata.version("click") < "8.2":
@@ -47,6 +47,14 @@ def split_opt(option: str) -> tuple[str, str]:
     return cast("Callable[[str], tuple[str, str]]", _split_opt)(option)
 
 
+def get_command(cli: click.Group, name: str) -> click.Command:
+    """Get a nested command by name."""
+    command: click.Command = cli
+    for part in name.split():
+        command = cast("click.Group", command).commands[part]
+    return command
+
+
 class EnumChoice(Choice):
     """An enum-based choice type.
 

src/_pytask/lock.py

@@ -43,6 +43,27 @@
     from _pytask.lockfile import LockfileState
 
 
+def _add_lock_command_options(
+    *, dry_run_help: str
+) -> Callable[[Callable[..., None]], Callable[..., None]]:
+    def decorator(func: Callable[..., None]) -> Callable[..., None]:
+        func = click.option(
+            "--dry-run",
+            is_flag=True,
+            default=False,
+            help=dry_run_help,
+        )(func)
+        return click.option(
+            "-y",
+            "--yes",
+            is_flag=True,
+            default=False,
+            help="Apply the changes without prompting for confirmation.",
+        )(func)
+
+    return decorator
+
+
 @dataclass(slots=True)
 class _PlannedChange:
     kind: str
@@ -237,6 +258,7 @@ def _run_lock_command(
     planner: Callable[[Session], list[_PlannedChange]] | None = None,
     planner_with_tasks: Callable[[Session, list[PTask]], list[_PlannedChange]]
     | None = None,
+    select_tasks: Callable[[Session], list[PTask]] | None = None,
     empty_message: str,
 ) -> int:
     _validate_confirmation_options(raw_config)
@@ -259,11 +281,8 @@ def _run_lock_command(
             session.dag = create_dag(session=session)
 
             if planner_with_tasks is not None:
-                tasks = (
-                    _select_tasks_with_ancestors(session)
-                    if raw_config["subcommand"] == "accept"
-                    else _select_tasks_exact(session)
-                )
+                assert select_tasks is not None
+                tasks = select_tasks(session)
                 planned_changes = planner_with_tasks(session, tasks)
             else:
                 assert planner is not None
@@ -305,70 +324,40 @@ def lock() -> None:
 
 
 @lock.command(cls=ColoredCommand)
-@click.option(
-    "--dry-run",
-    is_flag=True,
-    default=False,
-    help="Show which recorded states would be updated without writing changes.",
-)
-@click.option(
-    "-y",
-    "--yes",
-    is_flag=True,
-    default=False,
-    help="Apply the changes without prompting for confirmation.",
+@_add_lock_command_options(
+    dry_run_help="Show which recorded states would be updated without writing changes."
 )
 def accept(**raw_config: Any) -> None:
     """Accept the current state for selected tasks and their ancestors."""
-    raw_config["subcommand"] = "accept"
     sys.exit(
         _run_lock_command(
             raw_config,
             planner_with_tasks=_plan_accept_changes,
+            select_tasks=_select_tasks_with_ancestors,
             empty_message="No lockfile entries need updating.",
         )
     )
 
 
 @lock.command(cls=ColoredCommand)
-@click.option(
-    "--dry-run",
-    is_flag=True,
-    default=False,
-    help="Show which recorded states would be removed without writing changes.",
-)
-@click.option(
-    "-y",
-    "--yes",
-    is_flag=True,
-    default=False,
-    help="Apply the changes without prompting for confirmation.",
+@_add_lock_command_options(
+    dry_run_help="Show which recorded states would be removed without writing changes."
 )
 def reset(**raw_config: Any) -> None:
     """Remove recorded state for selected tasks."""
-    raw_config["subcommand"] = "reset"
     sys.exit(
         _run_lock_command(
             raw_config,
             planner_with_tasks=_plan_reset_changes,
+            select_tasks=_select_tasks_exact,
             empty_message="No lockfile entries need removing.",
         )
     )
 
 
 @lock.command(cls=ColoredCommand)
-@click.option(
-    "--dry-run",
-    is_flag=True,
-    default=False,
-    help="Show which stale entries would be removed without writing changes.",
-)
-@click.option(
-    "-y",
-    "--yes",
-    is_flag=True,
-    default=False,
-    help="Apply the changes without prompting for confirmation.",
+@_add_lock_command_options(
+    dry_run_help="Show which stale entries would be removed without writing changes."
 )
 def clean(**raw_config: Any) -> None:
     """Remove stale lockfile entries which no longer correspond to collected tasks."""

src/_pytask/mark/__init__.py

@@ -6,12 +6,12 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 from typing import Any
-from typing import cast
 
 import click
 from rich.table import Table
 
 from _pytask.click import ColoredCommand
+from _pytask.click import get_command
 from _pytask.console import console
 from _pytask.dag_utils import task_and_preceding_tasks
 from _pytask.exceptions import ConfigurationError
@@ -75,13 +75,6 @@ def markers(**raw_config: Any) -> NoReturn:
     sys.exit(session.exit_code)
 
 
-def _get_command(cli: click.Group, name: str) -> click.Command:
-    command: click.Command = cli
-    for part in name.split():
-        command = cast("click.Group", command).commands[part]
-    return command
-
-
 @hookimpl
 def pytask_extend_command_line_interface(cli: click.Group) -> None:
     """Add marker related options."""
@@ -110,7 +103,7 @@ def pytask_extend_command_line_interface(cli: click.Group) -> None:
         ),
     ]
     for command in ("build", "clean", "collect", "lock accept", "lock reset"):
-        target = _get_command(cli, command)
+        target = get_command(cli, command)
         target.params.extend(additional_build_parameters)
 
 

src/_pytask/parameters.py

@@ -5,14 +5,14 @@
 import importlib.util
 from pathlib import Path
 from typing import TYPE_CHECKING
-from typing import cast
 
 import click
 from click import Context
 from sqlalchemy.engine import URL
 from sqlalchemy.engine import make_url
 from sqlalchemy.exc import ArgumentError
 
+from _pytask.click import get_command
 from _pytask.config_utils import set_defaults_from_config
 from _pytask.path import import_path
 from _pytask.pluginmanager import hookimpl
@@ -180,13 +180,6 @@ def pytask_add_hooks(pm: PluginManager) -> None:
 )
 
 
-def _get_command(cli: click.Group, name: str) -> click.Command:
-    command: click.Command = cli
-    for part in name.split():
-        command = cast("click.Group", command).commands[part]
-    return command
-
-
 @hookimpl(trylast=True)
 def pytask_extend_command_line_interface(cli: click.Group) -> None:
     """Register general markers."""
@@ -203,7 +196,7 @@ def pytask_extend_command_line_interface(cli: click.Group) -> None:
         "markers",
         "profile",
     ):
-        target = _get_command(cli, command)
+        target = get_command(cli, command)
         target.params.extend((_CONFIG_OPTION, _HOOK_MODULE_OPTION, _PATH_ARGUMENT))
     for command in (
         "build",
@@ -214,7 +207,7 @@ def pytask_extend_command_line_interface(cli: click.Group) -> None:
         "lock reset",
         "profile",
     ):
-        target = _get_command(cli, command)
+        target = get_command(cli, command)
         target.params.extend([_IGNORE_OPTION, _EDITOR_URL_SCHEME_OPTION])
     for command in ("build",):
         cli.commands[command].params.append(_VERBOSE_OPTION)