diff options
Diffstat (limited to 'Lib')
-rw-r--r-- | Lib/test/test_tracemalloc.py | 20 | ||||
-rw-r--r-- | Lib/tracemalloc.py | 13 |
2 files changed, 33 insertions, 0 deletions
diff --git a/Lib/test/test_tracemalloc.py b/Lib/test/test_tracemalloc.py index 3d2333f..d1e5aef 100644 --- a/Lib/test/test_tracemalloc.py +++ b/Lib/test/test_tracemalloc.py @@ -510,6 +510,26 @@ class TestSnapshot(unittest.TestCase): self.assertEqual(traceback[:2], (traceback[0], traceback[1])) + def test_format_traceback(self): + snapshot, snapshot2 = create_snapshots() + def getline(filename, lineno): + return ' <%s, %s>' % (filename, lineno) + with unittest.mock.patch('tracemalloc.linecache.getline', + side_effect=getline): + tb = snapshot.traces[0].traceback + self.assertEqual(tb.format(), + [' File "a.py", line 2', + ' <a.py, 2>', + ' File "b.py", line 4', + ' <b.py, 4>']) + + self.assertEqual(tb.format(limit=1), + [' File "a.py", line 2', + ' <a.py, 2>']) + + self.assertEqual(tb.format(limit=-1), + []) + class TestFilters(unittest.TestCase): maxDiff = 2048 diff --git a/Lib/tracemalloc.py b/Lib/tracemalloc.py index b075946..6f0a234 100644 --- a/Lib/tracemalloc.py +++ b/Lib/tracemalloc.py @@ -1,6 +1,7 @@ from collections import Sequence from functools import total_ordering import fnmatch +import linecache import os.path import pickle @@ -205,6 +206,18 @@ class Traceback(Sequence): def __repr__(self): return "<Traceback %r>" % (tuple(self),) + def format(self, limit=None): + lines = [] + if limit is not None and limit < 0: + return lines + for frame in self[:limit]: + lines.append(' File "%s", line %s' + % (frame.filename, frame.lineno)) + line = linecache.getline(frame.filename, frame.lineno).strip() + if line: + lines.append(' %s' % line) + return lines + def get_object_traceback(obj): """ |