gh-139946: distinguish stdout or stderr when colorizing output in argparse (#140495)

Co-authored-by: blurb-it[bot] <43283697+blurb-it[bot]@users.noreply.github.com>
Co-authored-by: Savannah Ostrowski <savannah@python.org>
This commit is contained in:
Frost Ming 2025-12-08 12:08:06 +08:00 committed by GitHub
parent 3fa1425bfb
commit 7099af8f5e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 65 additions and 14 deletions

View file

@ -89,8 +89,8 @@
import os as _os
import re as _re
import sys as _sys
from gettext import gettext as _, ngettext
from gettext import gettext as _
from gettext import ngettext
SUPPRESS = '==SUPPRESS=='
@ -191,10 +191,10 @@ def __init__(
self._set_color(False)
def _set_color(self, color):
def _set_color(self, color, *, file=None):
from _colorize import can_colorize, decolor, get_theme
if color and can_colorize():
if color and can_colorize(file=file):
self._theme = get_theme(force_color=True).argparse
self._decolor = decolor
else:
@ -1675,7 +1675,7 @@ def _get_optional_kwargs(self, *args, **kwargs):
option_strings = []
for option_string in args:
# error on strings that don't start with an appropriate prefix
if not option_string[0] in self.prefix_chars:
if option_string[0] not in self.prefix_chars:
raise ValueError(
f'invalid option string {option_string!r}: '
f'must start with a character {self.prefix_chars!r}')
@ -2455,7 +2455,7 @@ def _parse_optional(self, arg_string):
return None
# if it doesn't start with a prefix, it was meant to be positional
if not arg_string[0] in self.prefix_chars:
if arg_string[0] not in self.prefix_chars:
return None
# if the option string is present in the parser, return the action
@ -2717,14 +2717,16 @@ def _check_value(self, action, value):
# Help-formatting methods
# =======================
def format_usage(self):
formatter = self._get_formatter()
def format_usage(self, formatter=None):
if formatter is None:
formatter = self._get_formatter()
formatter.add_usage(self.usage, self._actions,
self._mutually_exclusive_groups)
return formatter.format_help()
def format_help(self):
formatter = self._get_formatter()
def format_help(self, formatter=None):
if formatter is None:
formatter = self._get_formatter()
# usage
formatter.add_usage(self.usage, self._actions,
@ -2746,9 +2748,9 @@ def format_help(self):
# determine help from format above
return formatter.format_help()
def _get_formatter(self):
def _get_formatter(self, file=None):
formatter = self.formatter_class(prog=self.prog)
formatter._set_color(self.color)
formatter._set_color(self.color, file=file)
return formatter
def _get_validation_formatter(self):
@ -2765,12 +2767,26 @@ def _get_validation_formatter(self):
def print_usage(self, file=None):
if file is None:
file = _sys.stdout
self._print_message(self.format_usage(), file)
formatter = self._get_formatter(file=file)
try:
usage_text = self.format_usage(formatter=formatter)
except TypeError:
# Backward compatibility for formatter classes that
# do not accept the 'formatter' keyword argument.
usage_text = self.format_usage()
self._print_message(usage_text, file)
def print_help(self, file=None):
if file is None:
file = _sys.stdout
self._print_message(self.format_help(), file)
formatter = self._get_formatter(file=file)
try:
help_text = self.format_help(formatter=formatter)
except TypeError:
# Backward compatibility for formatter classes that
# do not accept the 'formatter' keyword argument.
help_text = self.format_help()
self._print_message(help_text, file)
def _print_message(self, message, file=None):
if message:

View file

@ -7558,6 +7558,40 @@ def test_error_and_warning_not_colorized_when_disabled(self):
self.assertNotIn('\x1b[', warn)
self.assertIn('warning:', warn)
def test_print_help_uses_target_file_for_color_decision(self):
parser = argparse.ArgumentParser(prog='PROG', color=True)
parser.add_argument('--opt')
output = io.StringIO()
calls = []
def fake_can_colorize(*, file=None):
calls.append(file)
return file is None
with swap_attr(_colorize, 'can_colorize', fake_can_colorize):
parser.print_help(file=output)
self.assertIs(calls[-1], output)
self.assertIn(output, calls)
self.assertNotIn('\x1b[', output.getvalue())
def test_print_usage_uses_target_file_for_color_decision(self):
parser = argparse.ArgumentParser(prog='PROG', color=True)
parser.add_argument('--opt')
output = io.StringIO()
calls = []
def fake_can_colorize(*, file=None):
calls.append(file)
return file is None
with swap_attr(_colorize, 'can_colorize', fake_can_colorize):
parser.print_usage(file=output)
self.assertIs(calls[-1], output)
self.assertIn(output, calls)
self.assertNotIn('\x1b[', output.getvalue())
class TestModule(unittest.TestCase):
def test_deprecated__version__(self):

View file

@ -0,0 +1 @@
Distinguish stdout and stderr when colorizing output in argparse module.