MNT: optionally collect gc in memleak.py

richardsheridan · richardsheridan · commit 15868033bf23 · 2023-03-19T03:33:10.000-04:00
diff --git a/tools/memleak.py b/tools/memleak.py
@@ -12,7 +12,7 @@
 import numpy as np
 
 
-def run_memleak_test(bench, iterations, report):
+def run_memleak_test(bench, iterations, report, collect):
     tracemalloc.start()
 
     starti = min(50, iterations // 2)
@@ -31,7 +31,8 @@ def run_memleak_test(bench, iterations, report):
     for i in range(endi):
         bench()
 
-        gc.collect()
+        if collect:
+            gc.collect()
 
         rss = p.memory_info().rss
         malloc, peak = tracemalloc.get_traced_memory()
@@ -135,6 +136,8 @@ def __call__(self):
     parser.add_argument('--interactive', action='store_true',
                         help="Turn on interactive mode to actually open "
                         "windows.  Only works with some GUI backends.")
+    parser.add_argument('--gc-collect', action='store_true',
+                        help="Run `gc.collect()` after every iteration")
 
     args = parser.parse_args()
 
@@ -146,4 +149,4 @@ def __call__(self):
         plt.ion()
 
     run_memleak_test(
-        MemleakTest(args.empty), args.iterations[0], args.report[0])
+        MemleakTest(args.empty), args.iterations[0], args.report[0], args.gc_collect)
