Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions Doc/whatsnew/3.15.rst
Original file line number Diff line number Diff line change
Expand Up @@ -1070,6 +1070,10 @@ tarfile
timeit
------

* The output of the :mod:`timeit` command-line interface is colored by default.
This can be controlled with
:ref:`environment variables <using-on-controlling-color>`.
(Contributed by Hugo van Kemenade in :gh:`146609`.)
* The command-line interface now colorizes error tracebacks
by default. This can be controlled with
:ref:`environment variables <using-on-controlling-color>`.
Expand Down
16 changes: 16 additions & 0 deletions Lib/_colorize.py
Original file line number Diff line number Diff line change
Expand Up @@ -323,6 +323,18 @@ 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_best: str = ANSIColors.GREEN
reset: str = ANSIColors.RESET


@dataclass(frozen=True, kw_only=True)
class Traceback(ThemeSection):
type: str = ANSIColors.BOLD_MAGENTA
Expand Down Expand Up @@ -356,6 +368,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)

Expand All @@ -366,6 +379,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:
Expand All @@ -379,6 +393,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,
)
Expand All @@ -396,6 +411,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(),
)
Expand Down
42 changes: 38 additions & 4 deletions Lib/test/test_timeit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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'])
Expand Down Expand Up @@ -398,5 +402,35 @@ def callback(a, b):
self.assertEqual(s.getvalue(), expected)


if __name__ == '__main__':
class TestTimeitColor(unittest.TestCase):

fake_stmt = TestTimeit.fake_stmt
run_main = TestTimeit.run_main

@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()
43 changes: 28 additions & 15 deletions Lib/timeit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -362,24 +367,32 @@ 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)))

best = min(timings)
worst = max(timings)
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}"
)

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 ({format_time(worst)})"
f"{theme.warning} was more than four times slower than the "
f"{theme.warning_best}best time ({format_time(best)})"
f"{theme.warning}.{reset}",
UserWarning, "", 0,
)
return None


Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add colour to :mod:`timeit` CLI output. Patch by Hugo van Kemenade.
Loading