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`.)
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
------

View file

@ -4,8 +4,9 @@
import io
from textwrap import dedent
from test.support import captured_stdout
from test.support import captured_stderr
from test.support import (
captured_stdout, captured_stderr, force_not_colorized,
)
# timeit's default number of iterations.
DEFAULT_NUMBER = 1000000
@ -351,11 +352,13 @@ 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'])

View file

@ -133,7 +133,7 @@ def __init__(self, stmt="pass", setup="pass", timer=default_timer,
exec(code, global_ns, local_ns)
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.
Typical use:
@ -149,6 +149,11 @@ def print_exc(self, file=None):
The optional file argument directs where the traceback is
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
if self.src is not None:
@ -158,7 +163,8 @@ def print_exc(self, file=None):
dummy_src_name)
# 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):
"""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
and returns another timer function (used for unit testing).
"""
import getopt
if args is None:
args = sys.argv[1:]
import getopt
import _colorize
colorize = _colorize.can_colorize()
try:
opts, args = getopt.getopt(args, "n:u:s:r:pvh",
["number=", "setup=", "repeat=",
@ -326,7 +335,7 @@ def callback(number, time_taken):
try:
number, _ = t.autorange(callback)
except:
t.print_exc()
t.print_exc(colorize=colorize)
return 1
if verbose:
@ -335,7 +344,7 @@ def callback(number, time_taken):
try:
raw_timings = t.repeat(repeat, number)
except:
t.print_exc()
t.print_exc(colorize=colorize)
return 1
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)'."""
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):
"""Like print_exc() but return a string."""

View file

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