88import inspect
99import builtins
1010import unittest
11+ import unittest .mock
1112import re
1213import tempfile
1314import random
2425import json
2526import textwrap
2627import traceback
28+ import contextlib
2729from functools import partial
2830from pathlib import Path
2931
4143class TracebackCases (unittest .TestCase ):
4244# For now, a very minimal set of tests. I want to be sure that
4345# formatting of SyntaxErrors works based on changes for 2.1.
46+ def setUp (self ):
47+ super ().setUp ()
48+ self .colorize = traceback ._COLORIZE
49+ traceback ._COLORIZE = False
50+
51+ def tearDown (self ):
52+ super ().tearDown ()
53+ traceback ._COLORIZE = self .colorize
4454
4555def get_exception_format (self ,func ,exc ):
4656try :
@@ -521,7 +531,7 @@ def test_signatures(self):
521531self .assertEqual (
522532str (inspect .signature (traceback .print_exception )),
523533 ('(exc, /, value=<implicit>, tb=<implicit>, '
524- 'limit=None, file=None, chain=True)' ))
534+ 'limit=None, file=None, chain=True, **kwargs )' ))
525535
526536self .assertEqual (
527537str (inspect .signature (traceback .format_exception )),
@@ -3031,7 +3041,7 @@ def some_inner(k, v):
30313041
30323042def test_custom_format_frame (self ):
30333043class CustomStackSummary (traceback .StackSummary ):
3034- def format_frame_summary (self ,frame_summary ):
3044+ def format_frame_summary (self ,frame_summary , colorize = False ):
30353045return f'{ frame_summary .filename } { frame_summary .lineno }
30363046
30373047def some_inner ():
@@ -3056,7 +3066,7 @@ def g():
30563066tb = g ()
30573067
30583068class Skip_G (traceback .StackSummary ):
3059- def format_frame_summary (self ,frame_summary ):
3069+ def format_frame_summary (self ,frame_summary , colorize = False ):
30603070if frame_summary .name == 'g' :
30613071return None
30623072return super ().format_frame_summary (frame_summary )
@@ -3076,7 +3086,6 @@ def __repr__(self) -> str:
30763086raise Exception ("Unrepresentable" )
30773087
30783088class TestTracebackException (unittest .TestCase ):
3079-
30803089def do_test_smoke (self ,exc ,expected_type_str ):
30813090try :
30823091raise exc
@@ -4245,6 +4254,115 @@ def test_levenshtein_distance_short_circuit(self):
42454254res3 = traceback ._levenshtein_distance (a ,b ,threshold )
42464255self .assertGreater (res3 ,threshold ,msg = (a ,b ,threshold ))
42474256
4257+ class TestColorizedTraceback (unittest .TestCase ):
4258+ def test_colorized_traceback (self ):
4259+ def foo (* args ):
4260+ x = {'a' :{'b' :None }}
4261+ y = x ['a' ]['b' ]['c' ]
4262+
4263+ def baz (* args ):
4264+ return foo (1 ,2 ,3 ,4 )
4265+
4266+ def bar ():
4267+ return baz (1 ,
4268+ 2 ,3
4269+ ,4 )
4270+ try :
4271+ bar ()
4272+ except Exception as e :
4273+ exc = traceback .TracebackException .from_exception (
4274+ e ,capture_locals = True
4275+ )
4276+ lines = "" .join (exc .format (colorize = True ))
4277+ red = traceback ._ANSIColors .RED
4278+ boldr = traceback ._ANSIColors .BOLD_RED
4279+ reset = traceback ._ANSIColors .RESET
4280+ self .assertIn ("y = " + red + "x['a']['b']" + reset + boldr + "['c']" + reset ,lines )
4281+ self .assertIn ("return " + red + "foo" + reset + boldr + "(1,2,3,4)" + reset ,lines )
4282+ self .assertIn ("return " + red + "baz" + reset + boldr + "(1," + reset ,lines )
4283+ self .assertIn (boldr + "2,3" + reset ,lines )
4284+ self .assertIn (boldr + ",4)" + reset ,lines )
4285+ self .assertIn (red + "bar" + reset + boldr + "()" + reset ,lines )
4286+
4287+ def test_colorized_syntax_error (self ):
4288+ try :
4289+ compile ("a $ b" ,"<string>" ,"exec" )
4290+ except SyntaxError as e :
4291+ exc = traceback .TracebackException .from_exception (
4292+ e ,capture_locals = True
4293+ )
4294+ actual = "" .join (exc .format (colorize = True ))
4295+ red = traceback ._ANSIColors .RED
4296+ magenta = traceback ._ANSIColors .MAGENTA
4297+ boldm = traceback ._ANSIColors .BOLD_MAGENTA
4298+ boldr = traceback ._ANSIColors .BOLD_RED
4299+ reset = traceback ._ANSIColors .RESET
4300+ expected = "" .join ([
4301+ f' File{ magenta } { reset } { magenta } { reset } \n ' ,
4302+ f' a{ boldr } { reset } \n ' ,
4303+ f'{ boldr } { reset } \n ' ,
4304+ f'{ boldm } { reset } { magenta } { reset } \n ' ]
4305+ )
4306+ self .assertIn (expected ,actual )
4307+
4308+ def test_colorized_traceback_is_the_default (self ):
4309+ def foo ():
4310+ 1 / 0
4311+
4312+ from _testcapi import exception_print
4313+ try :
4314+ foo ()
4315+ self .fail ("No exception thrown." )
4316+ except Exception as e :
4317+ with captured_output ("stderr" )as tbstderr :
4318+ with unittest .mock .patch ('traceback._can_colorize' ,return_value = True ):
4319+ exception_print (e )
4320+ actual = tbstderr .getvalue ().splitlines ()
4321+
4322+ red = traceback ._ANSIColors .RED
4323+ boldr = traceback ._ANSIColors .BOLD_RED
4324+ magenta = traceback ._ANSIColors .MAGENTA
4325+ boldm = traceback ._ANSIColors .BOLD_MAGENTA
4326+ reset = traceback ._ANSIColors .RESET
4327+ lno_foo = foo .__code__ .co_firstlineno
4328+ expected = ['Traceback (most recent call last):' ,
4329+ f' File{ magenta } { __file__ } { reset }
4330+ f'line{ magenta } { lno_foo + 5 } { reset } { magenta } { reset } ,
4331+ f'{ red } { reset + boldr } { reset } ,
4332+ f'{ red } { reset + boldr } { reset } ,
4333+ f' File{ magenta } { __file__ } { reset }
4334+ f'line{ magenta } { lno_foo + 1 } { reset } { magenta } { reset } ,
4335+ f'{ red } { reset + boldr } { reset + red } { reset } ,
4336+ f'{ red } { reset + boldr } { reset + red } { reset } ,
4337+ f'{ boldm } { reset } { magenta } { reset } ]
4338+ self .assertEqual (actual ,expected )
4339+
4340+ def test_colorized_detection_checks_for_environment_variables (self ):
4341+ if sys .platform == "win32" :
4342+ virtual_patching = unittest .mock .patch ("nt._supports_virtual_terminal" ,return_value = True )
4343+ else :
4344+ virtual_patching = contextlib .nullcontext ()
4345+ with virtual_patching :
4346+ with unittest .mock .patch ("os.isatty" )as isatty_mock :
4347+ isatty_mock .return_value = True
4348+ with unittest .mock .patch ("os.environ" , {'TERM' :'dumb' }):
4349+ self .assertEqual (traceback ._can_colorize (),False )
4350+ with unittest .mock .patch ("os.environ" , {'PYTHON_COLORS' :'1' }):
4351+ self .assertEqual (traceback ._can_colorize (),True )
4352+ with unittest .mock .patch ("os.environ" , {'PYTHON_COLORS' :'0' }):
4353+ self .assertEqual (traceback ._can_colorize (),False )
4354+ with unittest .mock .patch ("os.environ" , {'NO_COLOR' :'1' }):
4355+ self .assertEqual (traceback ._can_colorize (),False )
4356+ with unittest .mock .patch ("os.environ" , {'NO_COLOR' :'1' ,"PYTHON_COLORS" :'1' }):
4357+ self .assertEqual (traceback ._can_colorize (),True )
4358+ with unittest .mock .patch ("os.environ" , {'FORCE_COLOR' :'1' }):
4359+ self .assertEqual (traceback ._can_colorize (),True )
4360+ with unittest .mock .patch ("os.environ" , {'FORCE_COLOR' :'1' ,'NO_COLOR' :'1' }):
4361+ self .assertEqual (traceback ._can_colorize (),False )
4362+ with unittest .mock .patch ("os.environ" , {'FORCE_COLOR' :'1' ,"PYTHON_COLORS" :'0' }):
4363+ self .assertEqual (traceback ._can_colorize (),False )
4364+ isatty_mock .return_value = False
4365+ self .assertEqual (traceback ._can_colorize (),False )
42484366
42494367if __name__ == "__main__" :
42504368unittest .main ()