gh-139374: colorize traceback when using timeit command-line interface (#139375)

---------

Signed-off-by: yihong0618 <zouzou0208@gmail.com>
Co-authored-by: Stan Ulbrych <89152624+StanFromIreland@users.noreply.github.com>
Co-authored-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>
Co-authored-by: Peter Bierma <zintensitydev@gmail.com>
Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>
This commit is contained in:
yihong 2025-09-28 19:49:18 +08:00 committed by GitHub
parent 666112376d
commit e18dda96c9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 31 additions and 9 deletions

View file

@ -536,6 +536,15 @@ tarfile
(Contributed by Christoph Walcher in :gh:`57911`.) (Contributed by Christoph Walcher in :gh:`57911`.)
timeit
------
* The command-line interface now colorizes error tracebacks
by default. This can be controlled with
:ref:`environment variables <using-on-controlling-color>`.
(Contributed by Yi Hong in :gh:`139374`.)
types types
------ ------

View file

@ -4,8 +4,9 @@
import io import io
from textwrap import dedent from textwrap import dedent
from test.support import captured_stdout from test.support import (
from test.support import captured_stderr captured_stdout, captured_stderr, force_not_colorized,
)
# timeit's default number of iterations. # timeit's default number of iterations.
DEFAULT_NUMBER = 1000000 DEFAULT_NUMBER = 1000000
@ -351,11 +352,13 @@ def test_main_with_time_unit(self):
self.assertEqual(error_stringio.getvalue(), self.assertEqual(error_stringio.getvalue(),
"Unrecognized unit. Please select nsec, usec, msec, or sec.\n") "Unrecognized unit. Please select nsec, usec, msec, or sec.\n")
@force_not_colorized
def test_main_exception(self): def test_main_exception(self):
with captured_stderr() as error_stringio: with captured_stderr() as error_stringio:
s = self.run_main(switches=['1/0']) s = self.run_main(switches=['1/0'])
self.assert_exc_string(error_stringio.getvalue(), 'ZeroDivisionError') self.assert_exc_string(error_stringio.getvalue(), 'ZeroDivisionError')
@force_not_colorized
def test_main_exception_fixed_reps(self): def test_main_exception_fixed_reps(self):
with captured_stderr() as error_stringio: with captured_stderr() as error_stringio:
s = self.run_main(switches=['-n1', '1/0']) s = self.run_main(switches=['-n1', '1/0'])

View file

@ -133,7 +133,7 @@ def __init__(self, stmt="pass", setup="pass", timer=default_timer,
exec(code, global_ns, local_ns) exec(code, global_ns, local_ns)
self.inner = local_ns["inner"] self.inner = local_ns["inner"]
def print_exc(self, file=None): def print_exc(self, file=None, **kwargs):
"""Helper to print a traceback from the timed code. """Helper to print a traceback from the timed code.
Typical use: Typical use:
@ -149,6 +149,11 @@ def print_exc(self, file=None):
The optional file argument directs where the traceback is The optional file argument directs where the traceback is
sent; it defaults to sys.stderr. sent; it defaults to sys.stderr.
The optional colorize keyword argument controls whether the
traceback is colorized; it defaults to False for programmatic
usage. When used from the command line, this is automatically
set based on terminal capabilities.
""" """
import linecache, traceback import linecache, traceback
if self.src is not None: if self.src is not None:
@ -158,7 +163,8 @@ def print_exc(self, file=None):
dummy_src_name) dummy_src_name)
# else the source is already stored somewhere else # else the source is already stored somewhere else
traceback.print_exc(file=file) kwargs['colorize'] = kwargs.get('colorize', False)
traceback.print_exc(file=file, **kwargs)
def timeit(self, number=default_number): def timeit(self, number=default_number):
"""Time 'number' executions of the main statement. """Time 'number' executions of the main statement.
@ -257,9 +263,12 @@ def main(args=None, *, _wrap_timer=None):
is not None, it must be a callable that accepts a timer function is not None, it must be a callable that accepts a timer function
and returns another timer function (used for unit testing). and returns another timer function (used for unit testing).
""" """
import getopt
if args is None: if args is None:
args = sys.argv[1:] args = sys.argv[1:]
import getopt import _colorize
colorize = _colorize.can_colorize()
try: try:
opts, args = getopt.getopt(args, "n:u:s:r:pvh", opts, args = getopt.getopt(args, "n:u:s:r:pvh",
["number=", "setup=", "repeat=", ["number=", "setup=", "repeat=",
@ -326,7 +335,7 @@ def callback(number, time_taken):
try: try:
number, _ = t.autorange(callback) number, _ = t.autorange(callback)
except: except:
t.print_exc() t.print_exc(colorize=colorize)
return 1 return 1
if verbose: if verbose:
@ -335,7 +344,7 @@ def callback(number, time_taken):
try: try:
raw_timings = t.repeat(repeat, number) raw_timings = t.repeat(repeat, number)
except: except:
t.print_exc() t.print_exc(colorize=colorize)
return 1 return 1
def format_time(dt): def format_time(dt):

View file

@ -206,9 +206,9 @@ def _safe_string(value, what, func=str):
# -- # --
def print_exc(limit=None, file=None, chain=True): def print_exc(limit=None, file=None, chain=True, **kwargs):
"""Shorthand for 'print_exception(sys.exception(), limit=limit, file=file, chain=chain)'.""" """Shorthand for 'print_exception(sys.exception(), limit=limit, file=file, chain=chain)'."""
print_exception(sys.exception(), limit=limit, file=file, chain=chain) print_exception(sys.exception(), limit=limit, file=file, chain=chain, **kwargs)
def format_exc(limit=None, chain=True): def format_exc(limit=None, chain=True):
"""Like print_exc() but return a string.""" """Like print_exc() but return a string."""

View file

@ -0,0 +1 @@
:mod:`timeit`: Add color to error tracebacks.