gh-130645: Add color to argparse help (GH-132323)

This commit is contained in:
Hugo van Kemenade 2025-05-02 16:06:10 +03:00 committed by GitHub
parent ba16ba3a18
commit 4701ff92d7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 342 additions and 31 deletions

View File

@ -74,7 +74,7 @@ ArgumentParser objects
prefix_chars='-', fromfile_prefix_chars=None, \
argument_default=None, conflict_handler='error', \
add_help=True, allow_abbrev=True, exit_on_error=True, \
suggest_on_error=False)
suggest_on_error=False, color=False)
Create a new :class:`ArgumentParser` object. All parameters should be passed
as keyword arguments. Each parameter has its own more detailed description
@ -111,7 +111,7 @@ ArgumentParser objects
* add_help_ - Add a ``-h/--help`` option to the parser (default: ``True``)
* allow_abbrev_ - Allows long options to be abbreviated if the
abbreviation is unambiguous. (default: ``True``)
abbreviation is unambiguous (default: ``True``)
* exit_on_error_ - Determines whether or not :class:`!ArgumentParser` exits with
error info when an error occurs. (default: ``True``)
@ -119,6 +119,7 @@ ArgumentParser objects
* suggest_on_error_ - Enables suggestions for mistyped argument choices
and subparser names (default: ``False``)
* color_ - Allow color output (default: ``False``)
.. versionchanged:: 3.5
*allow_abbrev* parameter was added.
@ -130,6 +131,9 @@ ArgumentParser objects
.. versionchanged:: 3.9
*exit_on_error* parameter was added.
.. versionchanged:: 3.14
*suggest_on_error* and *color* parameters were added.
The following sections describe how each of these are used.
@ -594,7 +598,8 @@ subparser names, the feature can be enabled by setting ``suggest_on_error`` to
``True``. Note that this only applies for arguments when the choices specified
are strings::
>>> parser = argparse.ArgumentParser(description='Process some integers.', suggest_on_error=True)
>>> parser = argparse.ArgumentParser(description='Process some integers.',
suggest_on_error=True)
>>> parser.add_argument('--action', choices=['sum', 'max'])
>>> parser.add_argument('integers', metavar='N', type=int, nargs='+',
... help='an integer for the accumulator')
@ -612,6 +617,33 @@ keyword argument::
.. versionadded:: 3.14
color
^^^^^
By default, the help message is printed in plain text. If you want to allow
color in help messages, you can enable it by setting ``color`` to ``True``::
>>> parser = argparse.ArgumentParser(description='Process some integers.',
... color=True)
>>> parser.add_argument('--action', choices=['sum', 'max'])
>>> parser.add_argument('integers', metavar='N', type=int, nargs='+',
... help='an integer for the accumulator')
>>> parser.parse_args(['--help'])
Even if a CLI author has enabled color, it can be
:ref:`controlled using environment variables <using-on-controlling-color>`.
If you're writing code that needs to be compatible with older Python versions
and want to opportunistically use ``color`` when it's available, you
can set it as an attribute after initializing the parser instead of using the
keyword argument::
>>> parser = argparse.ArgumentParser(description='Process some integers.')
>>> parser.color = True
.. versionadded:: next
The add_argument() method
-------------------------

View File

@ -17,6 +17,7 @@ class ANSIColors:
BLUE = "\x1b[34m"
CYAN = "\x1b[36m"
GREEN = "\x1b[32m"
GREY = "\x1b[90m"
MAGENTA = "\x1b[35m"
RED = "\x1b[31m"
WHITE = "\x1b[37m" # more like LIGHT GRAY
@ -60,10 +61,12 @@ class ANSIColors:
INTENSE_BACKGROUND_YELLOW = "\x1b[103m"
ColorCodes = set()
NoColors = ANSIColors()
for attr in dir(NoColors):
for attr, code in ANSIColors.__dict__.items():
if not attr.startswith("__"):
ColorCodes.add(code)
setattr(NoColors, attr, "")
@ -76,6 +79,13 @@ def get_colors(
return NoColors
def decolor(text: str) -> str:
"""Remove ANSI color codes from a string."""
for code in ColorCodes:
text = text.replace(code, "")
return text
def can_colorize(*, file: IO[str] | IO[bytes] | None = None) -> bool:
if file is None:
file = sys.stdout

View File

@ -161,18 +161,31 @@ class HelpFormatter(object):
provided by the class are considered an implementation detail.
"""
def __init__(self,
prog,
indent_increment=2,
max_help_position=24,
width=None):
def __init__(
self,
prog,
indent_increment=2,
max_help_position=24,
width=None,
prefix_chars='-',
color=False,
):
# default setting for width
if width is None:
import shutil
width = shutil.get_terminal_size().columns
width -= 2
from _colorize import ANSIColors, NoColors, can_colorize, decolor
if color and can_colorize():
self._ansi = ANSIColors()
self._decolor = decolor
else:
self._ansi = NoColors
self._decolor = lambda text: text
self._prefix_chars = prefix_chars
self._prog = prog
self._indent_increment = indent_increment
self._max_help_position = min(max_help_position,
@ -224,9 +237,15 @@ class HelpFormatter(object):
# add the heading if the section was non-empty
if self.heading is not SUPPRESS and self.heading is not None:
bold_blue = self.formatter._ansi.BOLD_BLUE
reset = self.formatter._ansi.RESET
current_indent = self.formatter._current_indent
heading_text = _('%(heading)s:') % dict(heading=self.heading)
heading = '%*s%s\n' % (current_indent, '', heading_text)
heading = (
f'{" " * current_indent}'
f'{bold_blue}{heading_text}{reset}\n'
)
else:
heading = ''
@ -295,16 +314,26 @@ class HelpFormatter(object):
if part and part is not SUPPRESS])
def _format_usage(self, usage, actions, groups, prefix):
bold_blue = self._ansi.BOLD_BLUE
bold_magenta = self._ansi.BOLD_MAGENTA
magenta = self._ansi.MAGENTA
reset = self._ansi.RESET
if prefix is None:
prefix = _('usage: ')
# if usage is specified, use that
if usage is not None:
usage = usage % dict(prog=self._prog)
usage = (
magenta
+ usage
% {"prog": f"{bold_magenta}{self._prog}{reset}{magenta}"}
+ reset
)
# if no optionals or positionals are available, usage is just prog
elif usage is None and not actions:
usage = '%(prog)s' % dict(prog=self._prog)
usage = f"{bold_magenta}{self._prog}{reset}"
# if optionals and positionals are available, calculate usage
elif usage is None:
@ -326,7 +355,7 @@ class HelpFormatter(object):
# wrap the usage parts if it's too long
text_width = self._width - self._current_indent
if len(prefix) + len(usage) > text_width:
if len(prefix) + len(self._decolor(usage)) > text_width:
# break usage into wrappable parts
opt_parts = self._get_actions_usage_parts(optionals, groups)
@ -342,12 +371,13 @@ class HelpFormatter(object):
else:
line_len = indent_length - 1
for part in parts:
if line_len + 1 + len(part) > text_width and line:
part_len = len(self._decolor(part))
if line_len + 1 + part_len > text_width and line:
lines.append(indent + ' '.join(line))
line = []
line_len = indent_length - 1
line.append(part)
line_len += len(part) + 1
line_len += part_len + 1
if line:
lines.append(indent + ' '.join(line))
if prefix is not None:
@ -355,8 +385,9 @@ class HelpFormatter(object):
return lines
# if prog is short, follow it with optionals or positionals
if len(prefix) + len(prog) <= 0.75 * text_width:
indent = ' ' * (len(prefix) + len(prog) + 1)
prog_len = len(self._decolor(prog))
if len(prefix) + prog_len <= 0.75 * text_width:
indent = ' ' * (len(prefix) + prog_len + 1)
if opt_parts:
lines = get_lines([prog] + opt_parts, indent, prefix)
lines.extend(get_lines(pos_parts, indent))
@ -379,12 +410,25 @@ class HelpFormatter(object):
# join lines into usage
usage = '\n'.join(lines)
usage = usage.removeprefix(prog)
usage = f"{bold_magenta}{prog}{reset}{usage}"
# prefix with 'usage:'
return '%s%s\n\n' % (prefix, usage)
return f'{bold_blue}{prefix}{reset}{usage}\n\n'
def _format_actions_usage(self, actions, groups):
return ' '.join(self._get_actions_usage_parts(actions, groups))
def _is_long_option(self, string):
return len(string) >= 2 and string[1] in self._prefix_chars
def _is_short_option(self, string):
return (
not self._is_long_option(string)
and len(string) >= 1
and string[0] in self._prefix_chars
)
def _get_actions_usage_parts(self, actions, groups):
# find group indices and identify actions in groups
group_actions = set()
@ -408,6 +452,10 @@ class HelpFormatter(object):
# collect all actions format strings
parts = []
cyan = self._ansi.CYAN
green = self._ansi.GREEN
yellow = self._ansi.YELLOW
reset = self._ansi.RESET
for action in actions:
# suppressed arguments are marked with None
@ -417,7 +465,7 @@ class HelpFormatter(object):
# produce all arg strings
elif not action.option_strings:
default = self._get_default_metavar_for_positional(action)
part = self._format_args(action, default)
part = green + self._format_args(action, default) + reset
# if it's in a group, strip the outer []
if action in group_actions:
@ -432,13 +480,21 @@ class HelpFormatter(object):
# -s or --long
if action.nargs == 0:
part = action.format_usage()
if self._is_long_option(part):
part = f"{cyan}{part}{reset}"
elif self._is_short_option(part):
part = f"{green}{part}{reset}"
# if the Optional takes a value, format is:
# -s ARGS or --long ARGS
else:
default = self._get_default_metavar_for_optional(action)
args_string = self._format_args(action, default)
part = '%s %s' % (option_string, args_string)
if self._is_long_option(option_string):
option_string = f"{cyan}{option_string}"
elif self._is_short_option(option_string):
option_string = f"{green}{option_string}"
part = f"{option_string} {yellow}{args_string}{reset}"
# make it look optional if it's not required or in a group
if not action.required and action not in group_actions:
@ -485,6 +541,7 @@ class HelpFormatter(object):
help_width = max(self._width - help_position, 11)
action_width = help_position - self._current_indent - 2
action_header = self._format_action_invocation(action)
action_header_no_color = self._decolor(action_header)
# no help; start on same line and add a final newline
if not action.help:
@ -492,9 +549,15 @@ class HelpFormatter(object):
action_header = '%*s%s\n' % tup
# short action name; start on the same line and pad two spaces
elif len(action_header) <= action_width:
tup = self._current_indent, '', action_width, action_header
elif len(action_header_no_color) <= action_width:
# calculate widths without color codes
action_header_color = action_header
tup = self._current_indent, '', action_width, action_header_no_color
action_header = '%*s%-*s ' % tup
# swap in the colored header
action_header = action_header.replace(
action_header_no_color, action_header_color
)
indent_first = 0
# long action name; start on the next line
@ -527,23 +590,47 @@ class HelpFormatter(object):
return self._join_parts(parts)
def _format_action_invocation(self, action):
bold_green = self._ansi.BOLD_GREEN
bold_cyan = self._ansi.BOLD_CYAN
bold_yellow = self._ansi.BOLD_YELLOW
reset = self._ansi.RESET
if not action.option_strings:
default = self._get_default_metavar_for_positional(action)
return ' '.join(self._metavar_formatter(action, default)(1))
return (
bold_green
+ ' '.join(self._metavar_formatter(action, default)(1))
+ reset
)
else:
def color_option_strings(strings):
parts = []
for s in strings:
if self._is_long_option(s):
parts.append(f"{bold_cyan}{s}{reset}")
elif self._is_short_option(s):
parts.append(f"{bold_green}{s}{reset}")
else:
parts.append(s)
return parts
# if the Optional doesn't take a value, format is:
# -s, --long
if action.nargs == 0:
return ', '.join(action.option_strings)
option_strings = color_option_strings(action.option_strings)
return ', '.join(option_strings)
# if the Optional takes a value, format is:
# -s, --long ARGS
else:
default = self._get_default_metavar_for_optional(action)
args_string = self._format_args(action, default)
return ', '.join(action.option_strings) + ' ' + args_string
option_strings = color_option_strings(action.option_strings)
args_string = (
f"{bold_yellow}{self._format_args(action, default)}{reset}"
)
return ', '.join(option_strings) + ' ' + args_string
def _metavar_formatter(self, action, default_metavar):
if action.metavar is not None:
@ -1157,6 +1244,7 @@ class _SubParsersAction(Action):
self._name_parser_map = {}
self._choices_actions = []
self._deprecated = set()
self._color = False
super(_SubParsersAction, self).__init__(
option_strings=option_strings,
@ -1172,6 +1260,10 @@ class _SubParsersAction(Action):
if kwargs.get('prog') is None:
kwargs['prog'] = '%s %s' % (self._prog_prefix, name)
# set color
if kwargs.get('color') is None:
kwargs['color'] = self._color
aliases = kwargs.pop('aliases', ())
if name in self._name_parser_map:
@ -1776,7 +1868,8 @@ class ArgumentParser(_AttributeHolder, _ActionsContainer):
- exit_on_error -- Determines whether or not ArgumentParser exits with
error info when an error occurs
- suggest_on_error - Enables suggestions for mistyped argument choices
and subparser names. (default: ``False``)
and subparser names (default: ``False``)
- color - Allow color output in help messages (default: ``False``)
"""
def __init__(self,
@ -1793,8 +1886,10 @@ class ArgumentParser(_AttributeHolder, _ActionsContainer):
add_help=True,
allow_abbrev=True,
exit_on_error=True,
suggest_on_error=False):
suggest_on_error=False,
*,
color=False,
):
superinit = super(ArgumentParser, self).__init__
superinit(description=description,
prefix_chars=prefix_chars,
@ -1810,6 +1905,7 @@ class ArgumentParser(_AttributeHolder, _ActionsContainer):
self.allow_abbrev = allow_abbrev
self.exit_on_error = exit_on_error
self.suggest_on_error = suggest_on_error
self.color = color
add_group = self.add_argument_group
self._positionals = add_group(_('positional arguments'))
@ -1881,6 +1977,7 @@ class ArgumentParser(_AttributeHolder, _ActionsContainer):
# create the parsers action and add it to the positionals list
parsers_class = self._pop_action_class(kwargs, 'parsers')
action = parsers_class(option_strings=[], **kwargs)
action._color = self.color
self._check_help(action)
self._subparsers._add_action(action)
@ -2630,7 +2727,16 @@ class ArgumentParser(_AttributeHolder, _ActionsContainer):
return formatter.format_help()
def _get_formatter(self):
return self.formatter_class(prog=self.prog)
if isinstance(self.formatter_class, type) and issubclass(
self.formatter_class, HelpFormatter
):
return self.formatter_class(
prog=self.prog,
prefix_chars=self.prefix_chars,
color=self.color,
)
else:
return self.formatter_class(prog=self.prog)
# =====================
# Help-printing methods

View File

@ -1,5 +1,6 @@
# Author: Steven J. Bethard <steven.bethard@gmail.com>.
import _colorize
import contextlib
import functools
import inspect
@ -7046,6 +7047,167 @@ class TestTranslations(TestTranslationsBase):
self.assertMsgidsEqual(argparse)
# ===========
# Color tests
# ===========
class TestColorized(TestCase):
def setUp(self):
super().setUp()
# Ensure color even if ran with NO_COLOR=1
_colorize.can_colorize = lambda *args, **kwargs: True
self.ansi = _colorize.ANSIColors()
def test_argparse_color(self):
# Arrange: create a parser with a bit of everything
parser = argparse.ArgumentParser(
color=True,
description="Colorful help",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
prefix_chars="-+",
prog="PROG",
)
group = parser.add_mutually_exclusive_group()
group.add_argument(
"-v", "--verbose", action="store_true", help="more spam"
)
group.add_argument(
"-q", "--quiet", action="store_true", help="less spam"
)
parser.add_argument("x", type=int, help="the base")
parser.add_argument(
"y", type=int, help="the exponent", deprecated=True
)
parser.add_argument(
"this_indeed_is_a_very_long_action_name",
type=int,
help="the exponent",
)
parser.add_argument(
"-o", "--optional1", action="store_true", deprecated=True
)
parser.add_argument("--optional2", help="pick one")
parser.add_argument("--optional3", choices=("X", "Y", "Z"))
parser.add_argument(
"--optional4", choices=("X", "Y", "Z"), help="pick one"
)
parser.add_argument(
"--optional5", choices=("X", "Y", "Z"), help="pick one"
)
parser.add_argument(
"--optional6", choices=("X", "Y", "Z"), help="pick one"
)
parser.add_argument(
"-p",
"--optional7",
choices=("Aaaaa", "Bbbbb", "Ccccc", "Ddddd"),
help="pick one",
)
parser.add_argument("+f")
parser.add_argument("++bar")
parser.add_argument("-+baz")
parser.add_argument("-c", "--count")
subparsers = parser.add_subparsers(
title="subcommands",
description="valid subcommands",
help="additional help",
)
subparsers.add_parser("sub1", deprecated=True, help="sub1 help")
sub2 = subparsers.add_parser("sub2", deprecated=True, help="sub2 help")
sub2.add_argument("--baz", choices=("X", "Y", "Z"), help="baz help")
heading = self.ansi.BOLD_BLUE
label, label_b = self.ansi.YELLOW, self.ansi.BOLD_YELLOW
long, long_b = self.ansi.CYAN, self.ansi.BOLD_CYAN
pos, pos_b = short, short_b = self.ansi.GREEN, self.ansi.BOLD_GREEN
sub = self.ansi.BOLD_GREEN
prog = self.ansi.BOLD_MAGENTA
reset = self.ansi.RESET
# Act
help_text = parser.format_help()
# Assert
self.assertEqual(
help_text,
textwrap.dedent(
f"""\
{heading}usage: {reset}{prog}PROG{reset} [{short}-h{reset}] [{short}-v{reset} | {short}-q{reset}] [{short}-o{reset}] [{long}--optional2 {label}OPTIONAL2{reset}] [{long}--optional3 {label}{{X,Y,Z}}{reset}]
[{long}--optional4 {label}{{X,Y,Z}}{reset}] [{long}--optional5 {label}{{X,Y,Z}}{reset}] [{long}--optional6 {label}{{X,Y,Z}}{reset}]
[{short}-p {label}{{Aaaaa,Bbbbb,Ccccc,Ddddd}}{reset}] [{short}+f {label}F{reset}] [{long}++bar {label}BAR{reset}] [{long}-+baz {label}BAZ{reset}]
[{short}-c {label}COUNT{reset}]
{pos}x{reset} {pos}y{reset} {pos}this_indeed_is_a_very_long_action_name{reset} {pos}{{sub1,sub2}} ...{reset}
Colorful help
{heading}positional arguments:{reset}
{pos_b}x{reset} the base
{pos_b}y{reset} the exponent
{pos_b}this_indeed_is_a_very_long_action_name{reset}
the exponent
{heading}options:{reset}
{short_b}-h{reset}, {long_b}--help{reset} show this help message and exit
{short_b}-v{reset}, {long_b}--verbose{reset} more spam (default: False)
{short_b}-q{reset}, {long_b}--quiet{reset} less spam (default: False)
{short_b}-o{reset}, {long_b}--optional1{reset}
{long_b}--optional2{reset} {label_b}OPTIONAL2{reset}
pick one (default: None)
{long_b}--optional3{reset} {label_b}{{X,Y,Z}}{reset}
{long_b}--optional4{reset} {label_b}{{X,Y,Z}}{reset} pick one (default: None)
{long_b}--optional5{reset} {label_b}{{X,Y,Z}}{reset} pick one (default: None)
{long_b}--optional6{reset} {label_b}{{X,Y,Z}}{reset} pick one (default: None)
{short_b}-p{reset}, {long_b}--optional7{reset} {label_b}{{Aaaaa,Bbbbb,Ccccc,Ddddd}}{reset}
pick one (default: None)
{short_b}+f{reset} {label_b}F{reset}
{long_b}++bar{reset} {label_b}BAR{reset}
{long_b}-+baz{reset} {label_b}BAZ{reset}
{short_b}-c{reset}, {long_b}--count{reset} {label_b}COUNT{reset}
{heading}subcommands:{reset}
valid subcommands
{sub}{{sub1,sub2}}{reset} additional help
{sub}sub1{reset} sub1 help
{sub}sub2{reset} sub2 help
"""
),
)
def test_argparse_color_usage(self):
# Arrange
parser = argparse.ArgumentParser(
add_help=False,
color=True,
description="Test prog and usage colors",
prog="PROG",
usage="[prefix] %(prog)s [suffix]",
)
heading = self.ansi.BOLD_BLUE
prog = self.ansi.BOLD_MAGENTA
reset = self.ansi.RESET
usage = self.ansi.MAGENTA
# Act
help_text = parser.format_help()
# Assert
self.assertEqual(
help_text,
textwrap.dedent(
f"""\
{heading}usage: {reset}{usage}[prefix] {prog}PROG{reset}{usage} [suffix]{reset}
Test prog and usage colors
"""
),
)
def tearDownModule():
# Remove global references to avoid looking like we have refleaks.
RFile.seen = {}

View File

@ -0,0 +1 @@
Add colour to :mod:`argparse` help output. Patch by Hugo van Kemenade.