added max_iterations parameter to memory_usage function; added corresponding test case

Author: mmcneill

memory_profiler.py

@@ -238,7 +238,7 @@ def run(self):
 
 def memory_usage(proc=-1, interval=.1, timeout=None, timestamps=False,
                  include_children=False, multiprocess=False, max_usage=False,
-                 retval=False, stream=None, backend=None):
+                 retval=False, stream=None, backend=None, max_iterations=None):
     """
     Return the memory usage of a process or piece of code
 
@@ -307,6 +307,8 @@ def memory_usage(proc=-1, interval=.1, timeout=None, timestamps=False,
     else:
         # for a Python function wait until it finishes
         max_iter = float('inf')
+        if max_iterations is not None:
+            max_iter = max_iterations
 
     if callable(proc):
         proc = (proc, (), {})
@@ -320,7 +322,9 @@ def memory_usage(proc=-1, interval=.1, timeout=None, timestamps=False,
         else:
             raise ValueError
 
+        current_iter = 0
         while True:
+            current_iter += 1
             child_conn, parent_conn = Pipe()  # this will store MemTimer's results
             p = MemTimer(os.getpid(), interval, child_conn, backend,
                          timestamps=timestamps,
@@ -349,7 +353,8 @@ def memory_usage(proc=-1, interval=.1, timeout=None, timestamps=False,
                 raise
 
             p.join(5 * interval)
-            if n_measurements > 4 or interval < 1e-6:
+            
+            if (n_measurements > 4) or (current_iter == max_iter) or (interval < 1e-6):
                 break
             interval /= 10.
     elif isinstance(proc, subprocess.Popen):

test/test_memory_usage.py

@@ -1,4 +1,5 @@
 from memory_profiler import memory_usage
+import os
 
 
 def some_func(*args, **kwargs):
@@ -10,6 +11,20 @@ def test_memory_usage():
     mem, ret = memory_usage((some_func, (1, 2), dict(a=1)), retval=True)
     assert ret[0] == (1, 2)
     assert ret[1] == dict(a=1)
+    
+    
+def write_line(filepath):
+    with open(filepath, 'a') as the_file:
+        the_file.write('Testing\n')
+
+def test_max_iterations():
+    # Check that memory_usage works with max_iterations set (for python functions).
+    this_dir = os.path.dirname(os.path.realpath(__file__))
+    file = os.path.join(this_dir, 'temp_test_max_iterations_file.txt')
+    mem = memory_usage((write_line, (file, ), dict()), max_usage=True, max_iterations=1)
+    n_lines = sum(1 for line in open(file))
+    os.remove(file)
+    assert n_lines == 1
 
 
 def test_return_value_consistency():
@@ -27,4 +42,5 @@ def test_return_value_consistency():
 
 if __name__ == "__main__":
     test_memory_usage()
+    test_max_iterations()
     test_return_value_consistency()