bpo-40630: Add tracemalloc.reset_peak (GH-20102)

The reset_peak function sets the peak memory size to the current size,
representing a resetting of that metric. This allows for recording the
peak of specific sections of code, ignoring other code that may have
had a higher peak (since the most recent `tracemalloc.start()` or
tracemalloc.clear_traces()` call).

1 files changed, 24 insertions, 0 deletions

diff --git a/Lib/test/test_tracemalloc.py b/Lib/test/test_tracemalloc.py
index 635a9d3..c5ae4e6 100644
--- a/Lib/test/test_tracemalloc.py
+++ b/Lib/test/test_tracemalloc.py
@@ -246,6 +246,30 @@ class TestTracemallocEnabled(unittest.TestCase):
         traceback2 = tracemalloc.get_object_traceback(obj)
         self.assertIsNone(traceback2)
 
+    def test_reset_peak(self):
+        # Python allocates some internals objects, so the test must tolerate
+        # a small difference between the expected size and the real usage
+        tracemalloc.clear_traces()
+
+        # Example: allocate a large piece of memory, temporarily
+        large_sum = sum(list(range(100000)))
+        size1, peak1 = tracemalloc.get_traced_memory()
+
+        # reset_peak() resets peak to traced memory: peak2 < peak1
+        tracemalloc.reset_peak()
+        size2, peak2 = tracemalloc.get_traced_memory()
+        self.assertGreaterEqual(peak2, size2)
+        self.assertLess(peak2, peak1)
+
+        # check that peak continue to be updated if new memory is allocated:
+        # peak3 > peak2
+        obj_size = 1024 * 1024
+        obj, obj_traceback = allocate_bytes(obj_size)
+        size3, peak3 = tracemalloc.get_traced_memory()
+        self.assertGreaterEqual(peak3, size3)
+        self.assertGreater(peak3, peak2)
+        self.assertGreaterEqual(peak3 - peak2, obj_size)
+
     def test_is_tracing(self):
         tracemalloc.stop()
         self.assertFalse(tracemalloc.is_tracing())