Test: add non-regression test for handling of exceptions.

Also, make memory_profiler return in case of exception
instead of shutting down the full program.

Author: fabianp

Makefile

@@ -14,3 +14,4 @@ test:
 	$(PYTHON) test/test_import.py
 	$(PYTHON) test/test_memory_usage.py
 	$(PYTHON) test/test_precision_import.py
+	$(PYTHON) test/test_exception.py

README.rst

@@ -447,6 +447,8 @@ cleanup.
 
 `Benjamin Bengfort <https://github.com/bbengfort>`_ added support for tracking the usage of individual child processes and plotting them.
 
+`Muhammad Haseeb Tariq <https://github.com/mhaseebtariq>`_ fixed issue #152, which made the whole interpreter hang on functions that launched an exception.
+
 =========
  License
 =========

memory_profiler.py

@@ -326,6 +326,7 @@ def memory_usage(proc=-1, interval=.1, timeout=None, timestamps=False,
             raise ValueError
 
         while True:
+            exit_block = False
             child_conn, parent_conn = Pipe()  # this will store MemTimer's results
             p = MemTimer(os.getpid(), interval, child_conn, backend,
                          timestamps=timestamps,
@@ -338,23 +339,21 @@ def memory_usage(proc=-1, interval=.1, timeout=None, timestamps=False,
             # Therefore, the whole process hangs indefinitely. Here, we are ensuring that the process gets killed!
             try:
                 returned = f(*args, **kw)
+                parent_conn.send(0)  # finish timing
+                ret = parent_conn.recv()
+                n_measurements = parent_conn.recv()
+                if retval:
+                    ret = ret, returned
             except Exception:
-                sys.stderr.write(traceback.format_exc())
                 if has_psutil:
                     parent = psutil.Process(os.getpid())
                     for child in parent.children(recursive=True):
                         os.kill(child.pid, SIGKILL)
-                    os.kill(os.getpid(), SIGKILL)
-                else:
-                    sys.exit()
+                p.join(0)
+                raise
 
-            parent_conn.send(0)  # finish timing
-            ret = parent_conn.recv()
-            n_measurements = parent_conn.recv()
-            if retval:
-                ret = ret, returned
             p.join(5 * interval)
-            if n_measurements > 4 or interval < 1e-6:
+            if exit_block or n_measurements > 4 or interval < 1e-6:
                 break
             interval /= 10.
     elif isinstance(proc, subprocess.Popen):

test/test_exception.py

@@ -0,0 +1,12 @@
+
+# make sure that memory profiler does not hang on exception
+from memory_profiler import memory_usage
+
+def foo():
+    raise NotImplementedError('Error')
+
+try:
+    out = memory_usage((foo, tuple(), {}), timeout=1)
+except NotImplementedError:
+    pass
+print('Success')