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

diff --git a/Doc/library/tracemalloc.rst b/Doc/library/tracemalloc.rst
index e49d4ca4512c9991ec8543fdbce19e55d7102a5e..b79c5f63e707d31e0be4c69f97536ae32b85f6af 100644 (file)
--- a/Doc/library/tracemalloc.rst
+++ b/Doc/library/tracemalloc.rst
@@ -118,7 +118,6 @@ Get the traceback of a memory block
 
 Code to display the traceback of the biggest memory block::
 
-    import linecache
     import tracemalloc
 
     # Store 25 frames
@@ -132,12 +131,8 @@ Code to display the traceback of the biggest memory block::
     # pick the biggest memory block
     stat = top_stats[0]
     print("%s memory blocks: %.1f KiB" % (stat.count, stat.size / 1024))
-    for frame in stat.traceback:
-        print('  File "%s", line %s' % (frame.filename, frame.lineno))
-        line = linecache.getline(frame.filename, frame.lineno)
-        line = line.strip()
-        if line:
-            print('    ' + line)
+    for line in stat.traceback.format():
+        print(line)
 
 Example of output of the Python test suite (traceback limited to 25 frames)::
 
@@ -602,4 +597,26 @@ Traceback
    The :attr:`Trace.traceback` attribute is an instance of :class:`Traceback`
    instance.
 
+   .. method:: format(limit=None)
+
+      Format the traceback as a list of lines with newlines.  Use the
+      :mod:`linecache` module to retrieve lines from the source code.  If
+      *limit* is set, only format the *limit* most recent frames.
+
+      Similar to the :func:`traceback.format_tb` function, except that
+      :meth:`format` does not include newlines.
+
+      Example::
+
+          print("Traceback (most recent call first):")
+          for line in traceback:
+              print(line)
+
+      Output::
+
+          Traceback (most recent call first):
+            File "test.py", line 9
+              obj = Object()
+            File "test.py", line 12
+              tb = tracemalloc.get_object_traceback(f())
 

diff --git a/Lib/test/test_tracemalloc.py b/Lib/test/test_tracemalloc.py
index 3d2333fdd624c3a0ffe27414e2afba7ab9f04b4e..d1e5aef5c8cbee555badb2ef3101a56718ee73b3 100644 (file)
--- 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 b07594693974b4064c5be790eb9577dd0080a8f9..6f0a234244b4f83d1887ebc0e7dd75a06659df25 100644 (file)
--- 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):
     """

diff --git a/Misc/NEWS b/Misc/NEWS
index 4a717fc92bb1aaf0930cff95ba6d5820cf3d17da..48b2b89f399b7cc61d212bdd13fccdfe57f8661f 100644 (file)
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -23,6 +23,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #20616: Add a format() method to tracemalloc.Traceback.
+
 - Issue #19744: the ensurepip installation step now just prints a warning to
   stderr rather than failing outright if SSL/TLS is unavailable. This allows
   local installation of POSIX builds without SSL/TLS support.