Issue #20616: Add a format() method to tracemalloc.Traceback.

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):
     """