from __future__ import annotations import typing from collections.abc import Iterable from dataclasses import dataclass, replace from itertools import groupby from typing import Any, List, Sequence, Tuple, cast from rich.console import Console, NewLine from rich.markdown import Markdown from rich.padding import Padding from rich.table import Table from rich.text import Text from typing_extensions import Self, TypeAlias from cappa.arg import Arg, ArgAction, FinalArg from cappa.output import Displayable from cappa.subcommand import FinalSubcommand if typing.TYPE_CHECKING: from cappa.command import FinalCommand Dimension: TypeAlias = typing.Tuple[int, int, int, int] @dataclass(frozen=True) class FieldGroup: """A group of arguments or a subcommand sharing the same field_name. Either `args` or `subcommand` will be set, never both. """ field_name: str | None required: bool args: list[FinalArg[Any]] | None = None subcommand: FinalSubcommand | None = None @dataclass(frozen=True) class ArgGroup: """A group of arguments/subcommands organized by their Group.""" name: str field_groups: list[FieldGroup] @staticmethod def _by_group_key( arg: FinalArg[Any] | FinalSubcommand, ) -> tuple[int, int, str, bool]: return arg.group.key @staticmethod def _by_group(arg: FinalArg[Any] | FinalSubcommand) -> str: return arg.group.name @classmethod def collect( cls, command: FinalCommand[Any], include_hidden: bool = False ) -> list[ArgGroup]: """Collect and group arguments from a command by their Group.""" sorted_args = sorted(command.all_arguments, key=cls._by_group_key) result: list[ArgGroup] = [] for name, args in groupby(sorted_args, key=cls._by_group): items = [a for a in args if include_hidden or not a.hidden] field_grouped: dict[str | None, list[FinalArg[Any] | FinalSubcommand]] = {} for item in items: field_name = item.field_name or None if field_name not in field_grouped: field_grouped[field_name] = [] field_grouped[field_name].append(item) field_groups: list[FieldGroup] = [] for field_name, items in field_grouped.items(): # Validate that items are all Args or all Subcommands first_item = items[0] if isinstance(first_item, FinalArg): # Validate all are Args if not all(isinstance(item, FinalArg) for item in items): raise ValueError( # pragma: no cover f"FieldGroup with field_name={field_name!r} contains mixed types: " f"Args and Subcommands cannot share the same field_name" ) group_args = cast(List[FinalArg[Any]], items) # All items must have the same requiredness required_values = {arg.required for arg in group_args} assert len(required_values) == 1, ( f"All items with field_name={field_name!r} must have the same " f"required value, got: {required_values}" ) required = required_values.pop() field_groups.append( FieldGroup(field_name, required, args=group_args) ) else: assert len(items) == 1 assert isinstance(first_item, FinalSubcommand) field_groups.append( FieldGroup( field_name, first_item.required, subcommand=first_item, ) ) result.append(cls(name=name, field_groups=field_groups)) return result TextComponent = typing.Union[Text, Markdown, str] ArgFormat: TypeAlias = typing.Union[ TextComponent, typing.Callable[[Arg[Any]], typing.Union[TextComponent, None]], ] ArgFormats: TypeAlias = typing.Sequence[ArgFormat] class HelpFormattable(typing.Protocol): left_padding: Dimension arg_format: ArgFormat | ArgFormats default_format: str def long( self, command: FinalCommand[Any], prog: str ) -> list[Displayable]: ... # pragma: no cover def short( self, command: FinalCommand[Any], prog: str ) -> Displayable: ... # pragma: no cover @dataclass(frozen=True) class HelpFormatter(HelpFormattable): left_padding: Dimension = (0, 0, 0, 2) arg_format: ArgFormat | ArgFormats = ( Markdown("{help}"), Markdown("{choices}"), Markdown("{default}", style="dim italic"), ) default_format: str = "(Default: {default})" default: typing.ClassVar[HelpFormatter] def long(self, command: FinalCommand[Any], prog: str) -> list[Displayable]: arg_groups = ArgGroup.collect(command) lines: list[Displayable] = [] lines.append(add_short_args(prog, arg_groups)) if command.help: lines.append(NewLine()) lines.append(Padding(Markdown(f"**{command.help}**"), self.left_padding)) if command.description: lines.append(NewLine()) lines.append(Padding(Markdown(command.description), self.left_padding)) console = Console() lines.extend(add_long_args(console, self, arg_groups)) if command.epilog: lines.append(NewLine()) lines.append(Padding(Markdown(command.epilog), self.left_padding)) return lines def short(self, command: FinalCommand[Any], prog: str) -> Displayable: arg_groups = ArgGroup.collect(command) return add_short_args(prog, arg_groups) def with_arg_format( self, _format: ArgFormat | tuple[ArgFormat, ...], *formats: ArgFormat ) -> Self: format = cast( Tuple[ArgFormat, ...], _format if isinstance(_format, tuple) else (_format,) ) arg_format = (*format, *formats) return replace(self, arg_format=arg_format) def with_default_format(self, format: str) -> Self: return replace(self, default_format=format) HelpFormatter.default = HelpFormatter() def add_long_args( console: Console, help_formatter: HelpFormatter, arg_groups: list[ArgGroup] ) -> list[Table]: table = Table(box=None, expand=False, padding=help_formatter.left_padding) table.add_column(justify="left", ratio=1) table.add_column(style="cappa.help", ratio=2) for group in arg_groups: table.add_row( Text(group.name, style="cappa.group", justify="left"), Text(style="cappa.group"), ) for field_group in group.field_groups: if field_group.args: combined = format_arg_name(field_group, ", ") help_text = format_args(console, help_formatter, *field_group.args) table.add_row( Padding(combined, help_formatter.left_padding), help_text, ) else: assert field_group.subcommand subcommand = field_group.subcommand for option in subcommand.available_options(): table.add_row( *format_subcommand(help_formatter, option, subcommand) ) table.add_row() return [table] def format_args( console: Console, help_formatter: HelpFormattable, *args: FinalArg[Any] ) -> Displayable: """Format multiple args with the same field_name by concatenating their help texts.""" segments: list[TextComponent] = [] for arg in args: unknown_arg_format = help_formatter.arg_format if isinstance(unknown_arg_format, Iterable) and not isinstance( unknown_arg_format, str ): arg_format = cast(Sequence[ArgFormat], unknown_arg_format) else: arg_format = (unknown_arg_format,) for format_segment in arg_format: default = arg.show_default.format_default( arg.default, help_formatter.default_format ) choices = "" if arg.choices: choices = "Valid options: " + ", ".join(arg.choices) + "." context: dict[str, Any] = { "help": arg.help or "", "default": default, "choices": choices, "arg": arg, } if callable(format_segment): format_segment = cast(TextComponent, format_segment(arg)) format_segment_text = _get_text_component_text(format_segment) if not format_segment_text: continue formatted_text = format_segment_text.format(**context) segment = _replace_rich_text_component(format_segment, formatted_text) if segment: segments.append(segment) return _markdown_to_text(console, segments) def _markdown_to_text(console: Console, renderables: Sequence[TextComponent]) -> Text: # Render the text without line-wrapping because we're not yet printing it and it should reflow # until the point at which it's being printed to the user. render_console = Console(soft_wrap=True) render_options = render_console.options.update(no_wrap=True, overflow="ignore") result = Text() for renderable in renderables: if isinstance(renderable, Markdown): renderable_result = Text() for segment in render_console.render(renderable, render_options): text = segment.text if not text: continue if text.startswith(" "): # dedup leading spaces text = " " + text.lstrip() if text.endswith(" "): # dedup trailing spaces text = text.rstrip() + " " renderable_result.append(Text(text, style=segment.style or "", end="")) renderable_result.rstrip() if renderable_result: _append_to_text(result, renderable_result) else: if isinstance(renderable, str): renderable = Text.from_markup(renderable) _append_to_text(result, renderable) return result def _append_to_text(text: Text, component: Text) -> None: # When a component is multiline (like long, structured {help}), it looks weird when components # like {choices} or {default} are appended to the end of the last line. In such cases, add a newline. # Sort of a weird heuristic, we'll see if this holds well in practice. if text: text.append("\n" if "\n" in text.plain else " ") text.append(component) def _get_text_component_text(c: TextComponent) -> str: if isinstance(c, Text): return c.plain if isinstance(c, Markdown): return c.markup return c def _replace_rich_text_component(c: TextComponent, text: str) -> TextComponent: if isinstance(c, Text): return Text.from_markup( text, style=c.style, justify=c.justify, overflow=c.overflow, end=c.end, ) if isinstance(c, Markdown): return Markdown( text, code_theme=c.code_theme, justify=c.justify, style=c.style, hyperlinks=c.hyperlinks, inline_code_lexer=c.inline_code_lexer, inline_code_theme=c.inline_code_theme, ) return text def format_subcommand( help_formatter: HelpFormatter, command: FinalCommand[Any], subcommand: FinalSubcommand | None = None, ): canonical = command.real_name() parts = [f"[cappa.subcommand]{canonical}[/cappa.subcommand]"] if subcommand is not None: for alias in subcommand.visible_aliases_for(canonical): label = alias.name if alias.deprecated: label = f"{label} (deprecated)" parts.append(f"[cappa.subcommand]{label}[/cappa.subcommand]") return ( Padding(", ".join(parts), help_formatter.left_padding), command.help, ) def add_short_args(prog: str, arg_groups: list[ArgGroup]) -> str: segments: list[str] = [f"Usage: {prog}"] for group in arg_groups: for field_group in group.field_groups: segments.append(format_arg_name(field_group, ", ", n=1)) return " ".join(segments) def format_arg_name(item: FieldGroup, delimiter: str, *, n: int = 0) -> str: """Format the name(s) of args/subcommands for display. If given a FieldGroup, formats all items together with appropriate brackets. If given a single Arg/Subcommand, formats just that item. """ if item.args: parts: list[str] = [] for arg in item.args: num_args = arg.num_args has_value = ( not ArgAction.is_non_value_consuming(arg.action) and num_args.n != 0 ) arg_names = arg.names_str(delimiter, n=n) if not arg.is_option: arg_names = arg_names.upper() if num_args.n == -1: arg_names = f"{arg_names} ..." text = f"[cappa.arg]{arg_names}[/cappa.arg]" if arg.is_option and has_value: name = arg.value_name.upper() if num_args.n == -1: name = f"{name} ..." text = f"{text} [cappa.arg.name]{name}[/cappa.arg.name]" parts.append(text) combined = ", ".join(parts) if not item.required: return f"[{combined}]" return combined assert item.subcommand is not None arg_names = item.subcommand.names_str(",") return f"{{[cappa.subcommand]{arg_names}[/cappa.subcommand]}}" def format_subcommand_names(names: list[str]): return ", ".join(f"[cappa.subcommand]{a}[/cappa.subcommand]" for a in names)