Add colour to timeit CLI output

Author: hugovk

Lib/_colorize.py

@@ -323,6 +323,20 @@ class Syntax(ThemeSection):
     reset: str = ANSIColors.RESET
 
 
+@dataclass(frozen=True, kw_only=True)
+class Timeit(ThemeSection):
+    timing: str = ANSIColors.CYAN
+    best: str = ANSIColors.BOLD_GREEN
+    per_loop: str = ANSIColors.GREEN
+    arrow: str = ANSIColors.GREY
+    warning: str = ANSIColors.YELLOW
+    warning_worst: str = ANSIColors.MAGENTA
+    warning_worst_timing: str = ANSIColors.BOLD_MAGENTA
+    warning_best: str = ANSIColors.GREEN
+    warning_best_timing: str = ANSIColors.BOLD_GREEN
+    reset: str = ANSIColors.RESET
+
+
 @dataclass(frozen=True, kw_only=True)
 class Traceback(ThemeSection):
     type: str = ANSIColors.BOLD_MAGENTA
@@ -356,6 +370,7 @@ class Theme:
     difflib: Difflib = field(default_factory=Difflib)
     live_profiler: LiveProfiler = field(default_factory=LiveProfiler)
     syntax: Syntax = field(default_factory=Syntax)
+    timeit: Timeit = field(default_factory=Timeit)
     traceback: Traceback = field(default_factory=Traceback)
     unittest: Unittest = field(default_factory=Unittest)
 
@@ -366,6 +381,7 @@ def copy_with(
         difflib: Difflib | None = None,
         live_profiler: LiveProfiler | None = None,
         syntax: Syntax | None = None,
+        timeit: Timeit | None = None,
         traceback: Traceback | None = None,
         unittest: Unittest | None = None,
     ) -> Self:
@@ -379,6 +395,7 @@ def copy_with(
             difflib=difflib or self.difflib,
             live_profiler=live_profiler or self.live_profiler,
             syntax=syntax or self.syntax,
+            timeit=timeit or self.timeit,
             traceback=traceback or self.traceback,
             unittest=unittest or self.unittest,
         )
@@ -396,6 +413,7 @@ def no_colors(cls) -> Self:
             difflib=Difflib.no_colors(),
             live_profiler=LiveProfiler.no_colors(),
             syntax=Syntax.no_colors(),
+            timeit=Timeit.no_colors(),
             traceback=Traceback.no_colors(),
             unittest=Unittest.no_colors(),
         )

Lib/test/test_timeit.py

@@ -5,9 +5,14 @@
 from textwrap import dedent
 
 from test.support import (
-    captured_stdout, captured_stderr, force_not_colorized,
+    captured_stderr,
+    captured_stdout,
+    force_colorized,
+    force_not_colorized_test_class,
 )
 
+from _colorize import get_theme
+
 # timeit's default number of iterations.
 DEFAULT_NUMBER = 1000000
 
@@ -42,6 +47,7 @@ def wrap_timer(self, timer):
         self.saved_timer = timer
         return self
 
+@force_not_colorized_test_class
 class TestTimeit(unittest.TestCase):
 
     def tearDown(self):
@@ -352,13 +358,11 @@ def test_main_with_time_unit(self):
         self.assertEqual(error_stringio.getvalue(),
                     "Unrecognized unit. Please select nsec, usec, msec, or sec.\n")
 
-    @force_not_colorized
     def test_main_exception(self):
         with captured_stderr() as error_stringio:
             s = self.run_main(switches=['1/0'])
         self.assert_exc_string(error_stringio.getvalue(), 'ZeroDivisionError')
 
-    @force_not_colorized
     def test_main_exception_fixed_reps(self):
         with captured_stderr() as error_stringio:
             s = self.run_main(switches=['-n1', '1/0'])
@@ -398,5 +402,32 @@ def callback(a, b):
         self.assertEqual(s.getvalue(), expected)
 
 
-if __name__ == '__main__':
+class TestTimeitColor(TestTimeit):
+
+    @force_colorized
+    def test_main_colorized(self):
+        t = get_theme(force_color=True).timeit
+        s = self.run_main(seconds_per_increment=5.5)
+        self.assertEqual(
+            s,
+            "1 loop, best of 5: "
+            f"{t.best}5.5 sec {t.reset}"
+            f"{t.per_loop}per loop{t.reset}\n",
+        )
+
+    @force_colorized
+    def test_main_verbose_colorized(self):
+        t = get_theme(force_color=True).timeit
+        s = self.run_main(switches=["-v"])
+        self.assertEqual(
+            s,
+            f"1 loop {t.arrow}-> {t.timing}1 secs{t.reset}\n\n"
+            "raw times: "
+            f"{t.timing}1 sec, 1 sec, 1 sec, 1 sec, 1 sec{t.reset}\n\n"
+            f"1 loop, best of 5: {t.best}1 sec {t.reset}"
+            f"{t.per_loop}per loop{t.reset}\n",
+        )
+
+
+if __name__ == "__main__":
     unittest.main()

Lib/timeit.py

@@ -268,6 +268,8 @@ def main(args=None, *, _wrap_timer=None):
         args = sys.argv[1:]
     import _colorize
     colorize = _colorize.can_colorize()
+    theme = _colorize.get_theme(force_color=colorize).timeit
+    reset = theme.reset
 
     try:
         opts, args = getopt.getopt(args, "n:u:s:r:pvh",
@@ -328,10 +330,13 @@ def main(args=None, *, _wrap_timer=None):
         callback = None
         if verbose:
             def callback(number, time_taken):
-                msg = "{num} loop{s} -> {secs:.{prec}g} secs"
-                plural = (number != 1)
-                print(msg.format(num=number, s='s' if plural else '',
-                                 secs=time_taken, prec=precision))
+                s = "" if number == 1 else "s"
+                print(
+                    f"{number} loop{s} "
+                    f"{theme.arrow}-> "
+                    f"{theme.timing}{time_taken:.{precision}g} secs{reset}"
+                )
+
         try:
             number, _ = t.autorange(callback)
         except:
@@ -362,24 +367,38 @@ def format_time(dt):
         return "%.*g %s" % (precision, dt / scale, unit)
 
     if verbose:
-        print("raw times: %s" % ", ".join(map(format_time, raw_timings)))
+        raw = ", ".join(map(format_time, raw_timings))
+        print(f"raw times: {theme.timing}{raw}{reset}")
         print()
     timings = [dt / number for dt in raw_timings]
 
     best = min(timings)
-    print("%d loop%s, best of %d: %s per loop"
-          % (number, 's' if number != 1 else '',
-             repeat, format_time(best)))
+    s = "" if number == 1 else "s"
+    print(
+        f"{number} loop{s}, best of {repeat}: "
+        f"{theme.best}{format_time(best)} {reset}"
+        f"{theme.per_loop}per loop{reset}"
+    )
 
     best = min(timings)
     worst = max(timings)
     if worst >= best * 4:
         import warnings
-        warnings.warn_explicit("The test results are likely unreliable. "
-                               "The worst time (%s) was more than four times "
-                               "slower than the best time (%s)."
-                               % (format_time(worst), format_time(best)),
-                               UserWarning, '', 0)
+
+        print(file=sys.stderr)
+        warnings.warn_explicit(
+            f"{theme.warning}The test results are likely unreliable. "
+            f"The {theme.warning_worst}worst time ("
+            f"{theme.warning_worst_timing}{format_time(worst)}{reset}"
+            f"{theme.warning_worst})"
+            f"{theme.warning} was more than "
+            f"{theme.warning_worst}four times slower"
+            f"{theme.warning} than the "
+            f"{theme.warning_best}best time ("
+            f"{theme.warning_best_timing}{format_time(best)}{reset}"
+            f"{theme.warning_best}){theme.warning}.{reset}",
+            UserWarning, "", 0,
+        )
     return None