Merge pull request #224 from marcwil/master

Unpacking return value if max_usage is True

Author: fabianp

memory_profiler.py

@@ -337,6 +337,9 @@ def memory_usage(proc=-1, interval=.1, timeout=None, timestamps=False,
                 n_measurements = parent_conn.recv()
                 if retval:
                     ret = ret, returned
+                if max_usage:
+                    # Convert the one element list produced by MemTimer to a singular value
+                    ret = ret[0]
             except Exception:
                 parent = psutil.Process(os.getpid())
                 for child in parent.children(recursive=True):

test/test_memory_usage.py

@@ -11,5 +11,20 @@ def test_memory_usage():
     assert ret[0] == (1, 2)
     assert ret[1] == dict(a=1)
 
+
+def test_return_value_consistency():
+    # Test return values when watching process by PID
+    pid_mem_list = memory_usage(timeout=1)
+    assert type(pid_mem_list) == list, "Memory usage of process should be a list"
+    pid_mem_max = memory_usage(timeout=1, max_usage=True)
+    assert type(pid_mem_max) == float, "Max memory usage of process should be a number"
+    # Test return values when watching callable
+    func_mem_list = memory_usage((some_func, (42,), dict(a=42)))
+    assert type(func_mem_list) == list, "Memory usage of callable should be a list"
+    func_mem_max = memory_usage((some_func, (42,), dict(a=42)), max_usage=True)
+    assert type(func_mem_max) == float, "Max memory usage of callable should be a number"
+
+
 if __name__ == "__main__":
     test_memory_usage()
+    test_return_value_consistency()