tracemalloc: Fix slicing traces and fix slicing a traceback.

3 files changed, 18 insertions, 6 deletions

diff --git a/Lib/test/test_tracemalloc.py b/Lib/test/test_tracemalloc.py
index 484e313..3d2333f 100644
--- a/Lib/test/test_tracemalloc.py
+++ b/Lib/test/test_tracemalloc.py
@@ -123,7 +123,6 @@ class TestTracemallocEnabled(unittest.TestCase):
         self.assertEqual(len(traceback), 1)
         self.assertEqual(traceback, obj_traceback)
 
-
     def find_trace(self, traces, traceback):
         for trace in traces:
             if trace[1] == traceback._frames:
@@ -147,7 +146,6 @@ class TestTracemallocEnabled(unittest.TestCase):
         tracemalloc.stop()
         self.assertEqual(tracemalloc._get_traces(), [])
 
-
     def test_get_traces_intern_traceback(self):
         # dummy wrappers to get more useful and identical frames in the traceback
         def allocate_bytes2(size):
@@ -503,6 +501,14 @@ class TestSnapshot(unittest.TestCase):
         self.assertEqual(str(stat),
                          'a.py:5: size=5002 B (+5000 B), count=2 (+1), average=2501 B')
 
+    def test_slices(self):
+        snapshot, snapshot2 = create_snapshots()
+        self.assertEqual(snapshot.traces[:2],
+                         (snapshot.traces[0], snapshot.traces[1]))
+
+        traceback = snapshot.traces[0].traceback
+        self.assertEqual(traceback[:2],
+                         (traceback[0], traceback[1]))
 
 
 class TestFilters(unittest.TestCase):
diff --git a/Lib/tracemalloc.py b/Lib/tracemalloc.py
index 816f734..b075946 100644
--- a/Lib/tracemalloc.py
+++ b/Lib/tracemalloc.py
@@ -182,8 +182,10 @@ class Traceback(Sequence):
         return len(self._frames)
 
     def __getitem__(self, index):
-        trace = self._frames[index]
-        return Frame(trace)
+        if isinstance(index, slice):
+            return tuple(Frame(trace) for trace in self._frames[index])
+        else:
+            return Frame(self._frames[index])
 
     def __contains__(self, frame):
         return frame._frame in self._frames
@@ -259,8 +261,10 @@ class _Traces(Sequence):
         return len(self._traces)
 
     def __getitem__(self, index):
-        trace = self._traces[index]
-        return Trace(trace)
+        if isinstance(index, slice):
+            return tuple(Trace(trace) for trace in self._traces[index])
+        else:
+            return Trace(self._traces[index])
 
     def __contains__(self, trace):
         return trace._trace in self._traces
diff --git a/Misc/NEWS b/Misc/NEWS
index 89cec2e..a941da7 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -16,6 +16,8 @@ Core and Builtins
 Library
 -------
 
+- tracemalloc: Fix slicing traces and fix slicing a traceback.
+
 - Issue #20354: Fix an alignment issue in the tracemalloc module on 64-bit
   platforms. Bug seen on 64-bit Linux when using "make profile-opt".