Skip to content

Commit 67438ed

Browse files
ammaraskarammaraskar
authored
Fix traceback module (#9)
1 parent 44138dc commit 67438ed

File tree

6 files changed

+718
-637
lines changed

6 files changed

+718
-637
lines changed

‎Lib/test/test_traceback.py

Lines changed: 45 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,9 @@
1616

1717

1818
test_code = namedtuple('code', ['co_filename', 'co_name'])
19+
test_code.co_positions = lambda _: iter([(6, 6, 0, 0)])
1920
test_frame = namedtuple('frame', ['f_code', 'f_globals', 'f_locals'])
20-
test_tb = namedtuple('tb', ['tb_frame', 'tb_lineno', 'tb_next'])
21+
test_tb = namedtuple('tb', ['tb_frame', 'tb_lineno', 'tb_next', 'tb_lasti'])
2122

2223

2324
class TracebackCases(unittest.TestCase):
@@ -153,9 +154,9 @@ def do_test(firstlines, message, charset, lineno):
153154
self.assertTrue(stdout[2].endswith(err_line),
154155
"Invalid traceback line: {0!r} instead of {1!r}".format(
155156
stdout[2], err_line))
156-
self.assertTrue(stdout[3] == err_msg,
157+
self.assertTrue(stdout[4] == err_msg,
157158
"Invalid error message: {0!r} instead of {1!r}".format(
158-
stdout[3], err_msg))
159+
stdout[4], err_msg))
159160

160161
do_test("", "foo", "ascii", 3)
161162
for charset in ("ascii", "iso-8859-1", "utf-8", "GBK"):
@@ -299,9 +300,9 @@ def check_traceback_format(self, cleanup_func=None):
299300

300301
# Make sure that the traceback is properly indented.
301302
tb_lines = python_fmt.splitlines()
302-
self.assertEqual(len(tb_lines), 5)
303+
self.assertEqual(len(tb_lines), 7)
303304
banner = tb_lines[0]
304-
location, source_line = tb_lines[-2:]
305+
location, source_line = tb_lines[-3], tb_lines[-2]
305306
self.assertTrue(banner.startswith('Traceback'))
306307
self.assertTrue(location.startswith(' File'))
307308
self.assertTrue(source_line.startswith(' raise'))
@@ -365,12 +366,16 @@ def f():
365366
'Traceback (most recent call last):\n'
366367
f' File "{__file__}", line {lineno_f+5}, in _check_recursive_traceback_display\n'
367368
' f()\n'
369+
' ^^^\n'
368370
f' File "{__file__}", line {lineno_f+1}, in f\n'
369371
' f()\n'
372+
' ^^^\n'
370373
f' File "{__file__}", line {lineno_f+1}, in f\n'
371374
' f()\n'
375+
' ^^^\n'
372376
f' File "{__file__}", line {lineno_f+1}, in f\n'
373377
' f()\n'
378+
' ^^^\n'
374379
# XXX: The following line changes depending on whether the tests
375380
# are run through the interactive interpreter or with -m
376381
# It also varies depending on the platform (stack size)
@@ -411,19 +416,24 @@ def g(count=10):
411416
result_g = (
412417
f' File "{__file__}", line {lineno_g+2}, in g\n'
413418
' return g(count-1)\n'
419+
' ^^^^^^^^^^\n'
414420
f' File "{__file__}", line {lineno_g+2}, in g\n'
415421
' return g(count-1)\n'
422+
' ^^^^^^^^^^\n'
416423
f' File "{__file__}", line {lineno_g+2}, in g\n'
417424
' return g(count-1)\n'
425+
' ^^^^^^^^^^\n'
418426
' [Previous line repeated 7 more times]\n'
419427
f' File "{__file__}", line {lineno_g+3}, in g\n'
420428
' raise ValueError\n'
429+
' ^^^^^^^^^^^^^^^^\n'
421430
'ValueError\n'
422431
)
423432
tb_line = (
424433
'Traceback (most recent call last):\n'
425434
f' File "{__file__}", line {lineno_g+7}, in _check_recursive_traceback_display\n'
426435
' g()\n'
436+
' ^^^\n'
427437
)
428438
expected = (tb_line + result_g).splitlines()
429439
actual = stderr_g.getvalue().splitlines()
@@ -448,15 +458,20 @@ def h(count=10):
448458
'Traceback (most recent call last):\n'
449459
f' File "{__file__}", line {lineno_h+7}, in _check_recursive_traceback_display\n'
450460
' h()\n'
461+
' ^^^\n'
451462
f' File "{__file__}", line {lineno_h+2}, in h\n'
452463
' return h(count-1)\n'
464+
' ^^^^^^^^^^\n'
453465
f' File "{__file__}", line {lineno_h+2}, in h\n'
454466
' return h(count-1)\n'
467+
' ^^^^^^^^^^\n'
455468
f' File "{__file__}", line {lineno_h+2}, in h\n'
456469
' return h(count-1)\n'
470+
' ^^^^^^^^^^\n'
457471
' [Previous line repeated 7 more times]\n'
458472
f' File "{__file__}", line {lineno_h+3}, in h\n'
459473
' g()\n'
474+
' ^^^\n'
460475
)
461476
expected = (result_h + result_g).splitlines()
462477
actual = stderr_h.getvalue().splitlines()
@@ -473,18 +488,23 @@ def h(count=10):
473488
result_g = (
474489
f' File "{__file__}", line {lineno_g+2}, in g\n'
475490
' return g(count-1)\n'
491+
' ^^^^^^^^^^\n'
476492
f' File "{__file__}", line {lineno_g+2}, in g\n'
477493
' return g(count-1)\n'
494+
' ^^^^^^^^^^\n'
478495
f' File "{__file__}", line {lineno_g+2}, in g\n'
479496
' return g(count-1)\n'
497+
' ^^^^^^^^^^\n'
480498
f' File "{__file__}", line {lineno_g+3}, in g\n'
481499
' raise ValueError\n'
500+
' ^^^^^^^^^^^^^^^^\n'
482501
'ValueError\n'
483502
)
484503
tb_line = (
485504
'Traceback (most recent call last):\n'
486-
f' File "{__file__}", line {lineno_g+71}, in _check_recursive_traceback_display\n'
505+
f' File "{__file__}", line {lineno_g+81}, in _check_recursive_traceback_display\n'
487506
' g(traceback._RECURSIVE_CUTOFF)\n'
507+
' ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n'
488508
)
489509
expected = (tb_line + result_g).splitlines()
490510
actual = stderr_g.getvalue().splitlines()
@@ -501,19 +521,24 @@ def h(count=10):
501521
result_g = (
502522
f' File "{__file__}", line {lineno_g+2}, in g\n'
503523
' return g(count-1)\n'
524+
' ^^^^^^^^^^\n'
504525
f' File "{__file__}", line {lineno_g+2}, in g\n'
505526
' return g(count-1)\n'
527+
' ^^^^^^^^^^\n'
506528
f' File "{__file__}", line {lineno_g+2}, in g\n'
507529
' return g(count-1)\n'
530+
' ^^^^^^^^^^\n'
508531
' [Previous line repeated 1 more time]\n'
509532
f' File "{__file__}", line {lineno_g+3}, in g\n'
510533
' raise ValueError\n'
534+
' ^^^^^^^^^^^^^^^^\n'
511535
'ValueError\n'
512536
)
513537
tb_line = (
514538
'Traceback (most recent call last):\n'
515-
f' File "{__file__}", line {lineno_g+99}, in _check_recursive_traceback_display\n'
539+
f' File "{__file__}", line {lineno_g+114}, in _check_recursive_traceback_display\n'
516540
' g(traceback._RECURSIVE_CUTOFF + 1)\n'
541+
' ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n'
517542
)
518543
expected = (tb_line + result_g).splitlines()
519544
actual = stderr_g.getvalue().splitlines()
@@ -564,10 +589,10 @@ def __eq__(self, other):
564589
exception_print(exc_val)
565590

566591
tb = stderr_f.getvalue().strip().splitlines()
567-
self.assertEqual(11, len(tb))
568-
self.assertEqual(context_message.strip(), tb[5])
569-
self.assertIn('UnhashableException: ex2', tb[3])
570-
self.assertIn('UnhashableException: ex1', tb[10])
592+
self.assertEqual(13, len(tb))
593+
self.assertEqual(context_message.strip(), tb[6])
594+
self.assertIn('UnhashableException: ex2', tb[4])
595+
self.assertIn('UnhashableException: ex1', tb[12])
571596

572597

573598
cause_message = (
@@ -597,8 +622,8 @@ def zero_div(self):
597622

598623
def check_zero_div(self, msg):
599624
lines = msg.splitlines()
600-
self.assertTrue(lines[-3].startswith(' File'))
601-
self.assertIn('1/0 # In zero_div', lines[-2])
625+
self.assertTrue(lines[-4].startswith(' File'))
626+
self.assertIn('1/0 # In zero_div', lines[-3])
602627
self.assertTrue(lines[-1].startswith('ZeroDivisionError'), lines[-1])
603628

604629
def test_simple(self):
@@ -607,11 +632,11 @@ def test_simple(self):
607632
except ZeroDivisionError as _:
608633
e = _
609634
lines = self.get_report(e).splitlines()
610-
self.assertEqual(len(lines), 4)
635+
self.assertEqual(len(lines), 5)
611636
self.assertTrue(lines[0].startswith('Traceback'))
612637
self.assertTrue(lines[1].startswith(' File'))
613638
self.assertIn('1/0 # Marker', lines[2])
614-
self.assertTrue(lines[3].startswith('ZeroDivisionError'))
639+
self.assertTrue(lines[4].startswith('ZeroDivisionError'))
615640

616641
def test_cause(self):
617642
def inner_raise():
@@ -650,11 +675,11 @@ def test_context_suppression(self):
650675
except ZeroDivisionError as _:
651676
e = _
652677
lines = self.get_report(e).splitlines()
653-
self.assertEqual(len(lines), 4)
678+
self.assertEqual(len(lines), 5)
654679
self.assertTrue(lines[0].startswith('Traceback'))
655680
self.assertTrue(lines[1].startswith(' File'))
656681
self.assertIn('ZeroDivisionError from None', lines[2])
657-
self.assertTrue(lines[3].startswith('ZeroDivisionError'))
682+
self.assertTrue(lines[4].startswith('ZeroDivisionError'))
658683

659684
def test_cause_and_context(self):
660685
# When both a cause and a context are set, only the cause should be
@@ -1346,7 +1371,7 @@ def test_lookup_lines(self):
13461371
e = Exception("uh oh")
13471372
c = test_code('/foo.py', 'method')
13481373
f = test_frame(c, None, None)
1349-
tb = test_tb(f, 6, None)
1374+
tb = test_tb(f, 6, None, 0)
13501375
exc = traceback.TracebackException(Exception, e, tb, lookup_lines=False)
13511376
self.assertEqual(linecache.cache, {})
13521377
linecache.updatecache('/foo.py', globals())
@@ -1357,7 +1382,7 @@ def test_locals(self):
13571382
e = Exception("uh oh")
13581383
c = test_code('/foo.py', 'method')
13591384
f = test_frame(c, globals(), {'something': 1, 'other': 'string'})
1360-
tb = test_tb(f, 6, None)
1385+
tb = test_tb(f, 6, None, 0)
13611386
exc = traceback.TracebackException(
13621387
Exception, e, tb, capture_locals=True)
13631388
self.assertEqual(
@@ -1368,7 +1393,7 @@ def test_no_locals(self):
13681393
e = Exception("uh oh")
13691394
c = test_code('/foo.py', 'method')
13701395
f = test_frame(c, globals(), {'something': 1})
1371-
tb = test_tb(f, 6, None)
1396+
tb = test_tb(f, 6, None, 0)
13721397
exc = traceback.TracebackException(Exception, e, tb)
13731398
self.assertEqual(exc.stack[0].locals, None)
13741399

‎Lib/traceback.py

Lines changed: 60 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,8 @@ def extract_tb(tb, limit=None):
6969
trace. The line is a string with leading and trailing
7070
whitespace stripped; if the source is not available it is None.
7171
"""
72-
return StackSummary.extract(walk_tb(tb), limit=limit)
72+
return StackSummary._extract_from_extended_frame_gen(
73+
_walk_tb_with_full_positions(tb), limit=limit)
7374

7475
#
7576
# Exception formatting and output.
@@ -251,10 +252,12 @@ class FrameSummary:
251252
mapping the name to the repr() of the variable.
252253
"""
253254

254-
__slots__ = ('filename', 'lineno', 'name', '_line', 'locals')
255+
__slots__ = ('filename', 'lineno', 'end_lineno', 'colno', 'end_colno',
256+
'name', '_line', 'locals')
255257

256258
def __init__(self, filename, lineno, name, *, lookup_line=True,
257-
locals=None, line=None):
259+
locals=None, line=None,
260+
end_lineno=None, colno=None, end_colno=None):
258261
"""Construct a FrameSummary.
259262
260263
:param lookup_line: If True, `linecache` is consulted for the source
@@ -271,6 +274,9 @@ def __init__(self, filename, lineno, name, *, lookup_line=True,
271274
if lookup_line:
272275
self.line
273276
self.locals = {k: repr(v) for k, v in locals.items()} if locals else None
277+
self.end_lineno = end_lineno
278+
self.colno = colno
279+
self.end_colno = end_colno
274280

275281
def __eq__(self, other):
276282
if isinstance(other, FrameSummary):
@@ -295,11 +301,17 @@ def __repr__(self):
295301
def __len__(self):
296302
return 4
297303

304+
@property
305+
def _original_line(self):
306+
# Returns the line as-is from the source, without modifying whitespace.
307+
self.line
308+
return self._line
309+
298310
@property
299311
def line(self):
300312
if self._line is None:
301-
self._line = linecache.getline(self.filename, self.lineno).strip()
302-
return self._line
313+
self._line = linecache.getline(self.filename, self.lineno)
314+
return self._line.strip()
303315

304316

305317
def walk_stack(f):
@@ -309,7 +321,7 @@ def walk_stack(f):
309321
current stack is used. Usually used with StackSummary.extract.
310322
"""
311323
if f is None:
312-
f = sys._getframe().f_back.f_back
324+
f = sys._getframe().f_back.f_back.f_back.f_back
313325
while f is not None:
314326
yield f, f.f_lineno
315327
f = f.f_back
@@ -326,6 +338,19 @@ def walk_tb(tb):
326338
tb = tb.tb_next
327339

328340

341+
def _walk_tb_with_full_positions(tb):
342+
# Internal version of walk_tb that yields full code positions including
343+
# end line and column information.
344+
while tb is not None:
345+
yield tb.tb_frame, _get_code_position(tb.tb_frame.f_code, tb.tb_lasti)
346+
tb = tb.tb_next
347+
348+
349+
def _get_code_position(code, instruction_index):
350+
positions_gen = code.co_positions()
351+
return next(itertools.islice(positions_gen, instruction_index // 2, None))
352+
353+
329354
_RECURSIVE_CUTOFF = 3 # Also hardcoded in traceback.c.
330355

331356
class StackSummary(list):
@@ -345,6 +370,21 @@ def extract(klass, frame_gen, *, limit=None, lookup_lines=True,
345370
:param capture_locals: If True, the local variables from each frame will
346371
be captured as object representations into the FrameSummary.
347372
"""
373+
def extended_frame_gen():
374+
for f, lineno in frame_gen:
375+
yield f, (lineno, None, None, None)
376+
377+
return klass._extract_from_extended_frame_gen(
378+
extended_frame_gen(), limit=limit, lookup_lines=lookup_lines,
379+
capture_locals=capture_locals)
380+
381+
@classmethod
382+
def _extract_from_extended_frame_gen(klass, frame_gen, *, limit=None,
383+
lookup_lines=True, capture_locals=False):
384+
# Same as extract but operates on a frame generator that yields
385+
# (frame, (lineno, end_lineno, colno, end_colno)) in the stack.
386+
# Only lineno is required, the remaining fields can be empty if the
387+
# information is not available.
348388
if limit is None:
349389
limit = getattr(sys, 'tracebacklimit', None)
350390
if limit is not None and limit < 0:
@@ -357,7 +397,7 @@ def extract(klass, frame_gen, *, limit=None, lookup_lines=True,
357397

358398
result = klass()
359399
fnames = set()
360-
for f, lineno in frame_gen:
400+
for f, (lineno, end_lineno, colno, end_colno) in frame_gen:
361401
co = f.f_code
362402
filename = co.co_filename
363403
name = co.co_name
@@ -370,7 +410,8 @@ def extract(klass, frame_gen, *, limit=None, lookup_lines=True,
370410
else:
371411
f_locals = None
372412
result.append(FrameSummary(
373-
filename, lineno, name, lookup_line=False, locals=f_locals))
413+
filename, lineno, name, lookup_line=False, locals=f_locals,
414+
end_lineno=end_lineno, colno=colno, end_colno=end_colno))
374415
for filename in fnames:
375416
linecache.checkcache(filename)
376417
# If immediate lookup was desired, trigger lookups now.
@@ -437,6 +478,14 @@ def format(self):
437478
frame.filename, frame.lineno, frame.name))
438479
if frame.line:
439480
row.append(' {}\n'.format(frame.line.strip()))
481+
482+
stripped_characters = len(frame._original_line) - len(frame.line.lstrip())
483+
if frame.end_lineno == frame.lineno and frame.end_colno != 0:
484+
row.append(' ')
485+
row.append(' ' * (frame.colno - stripped_characters))
486+
row.append('^' * (frame.end_colno - frame.colno))
487+
row.append('\n')
488+
440489
if frame.locals:
441490
for name, value in sorted(frame.locals.items()):
442491
row.append(' {name} = {value}\n'.format(name=name, value=value))
@@ -491,8 +540,9 @@ def __init__(self, exc_type, exc_value, exc_traceback, *, limit=None,
491540
_seen.add(id(exc_value))
492541

493542
# TODO: locals.
494-
self.stack = StackSummary.extract(
495-
walk_tb(exc_traceback), limit=limit, lookup_lines=lookup_lines,
543+
self.stack = StackSummary._extract_from_extended_frame_gen(
544+
_walk_tb_with_full_positions(exc_traceback),
545+
limit=limit, lookup_lines=lookup_lines,
496546
capture_locals=capture_locals)
497547
self.exc_type = exc_type
498548
# Capture now to permit freeing resources: only complication is in the

0 commit comments

Comments
 (0)