Colorize error and warning messages in argparse

Author: savannahostrowski

Lib/_colorize.py

@@ -170,6 +170,8 @@ class Argparse(ThemeSection):
     label: str = ANSIColors.BOLD_YELLOW
     action: str = ANSIColors.BOLD_GREEN
     reset: str = ANSIColors.RESET
+    error: str = ANSIColors.BOLD_RED
+    warning: str = ANSIColors.BOLD_YELLOW
 
 
 @dataclass(frozen=True, kw_only=True)

Lib/argparse.py

@@ -2749,6 +2749,14 @@ def _print_message(self, message, file=None):
             except (AttributeError, OSError):
                 pass
 
+    def _get_theme(self, file=None):
+        from _colorize import can_colorize, get_theme
+
+        if self.color and can_colorize(file=file):
+            return get_theme(force_color=True).argparse
+        else:
+            return get_theme(force_no_color=True).argparse
+
     # ===============
     # Exiting methods
     # ===============
@@ -2768,13 +2776,19 @@ def error(self, message):
         should either exit or raise an exception.
         """
         self.print_usage(_sys.stderr)
+        theme = self._get_theme(file=_sys.stderr)
+        fmt = _('%(prog)s: error: %(message)s\n')
+        fmt = fmt.replace('error:', f'{theme.error}error:{theme.reset}')
+
         args = {'prog': self.prog, 'message': message}
-        self.exit(2, _('%(prog)s: error: %(message)s\n') % args)
+        self.exit(2, fmt % args)
 
     def _warning(self, message):
+        theme = self._get_theme(file=_sys.stderr)
+        fmt = _('%(prog)s: warning: %(message)s\n')
+        fmt = fmt.replace('warning:', f'{theme.warning}warning:{theme.reset}')
         args = {'prog': self.prog, 'message': message}
-        self._print_message(_('%(prog)s: warning: %(message)s\n') % args, _sys.stderr)
-
+        self._print_message(fmt % args, _sys.stderr)
 
 def __getattr__(name):
     if name == "__version__":

Lib/test/test_argparse.py

@@ -7370,6 +7370,45 @@ def test_subparser_prog_is_stored_without_color(self):
         help_text = demo_parser.format_help()
         self.assertNotIn('\x1b[', help_text)
 
+    def test_error_and_warning_keywords_colorized(self):
+        parser = argparse.ArgumentParser(prog='PROG')
+        parser.add_argument('foo')
+
+        with self.assertRaises(SystemExit):
+            with captured_stderr() as stderr:
+                parser.parse_args([])
+
+        err = stderr.getvalue()
+        error_color = self.theme.error
+        reset = self.theme.reset
+        self.assertIn(f'{error_color}error:{reset}', err)
+
+        with captured_stderr() as stderr:
+            parser._warning('test warning')
+
+        warn = stderr.getvalue()
+        warning_color = self.theme.warning
+        self.assertIn(f'{warning_color}warning:{reset}', warn)
+
+    def test_error_and_warning_not_colorized_when_disabled(self):
+        parser = argparse.ArgumentParser(prog='PROG', color=False)
+        parser.add_argument('foo')
+
+        with self.assertRaises(SystemExit):
+            with captured_stderr() as stderr:
+                parser.parse_args([])
+
+        err = stderr.getvalue()
+        self.assertNotIn('\x1b[', err)
+        self.assertIn('error:', err)
+
+        with captured_stderr() as stderr:
+            parser._warning('test warning')
+
+        warn = stderr.getvalue()
+        self.assertNotIn('\x1b[', warn)
+        self.assertIn('warning:', warn)
+
 
 class TestModule(unittest.TestCase):
     def test_deprecated__version__(self):