tracemalloc: Fix slicing traces and fix slicing a traceback.

2 files changed, 16 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