瀏覽代碼

Make sure an exception in any iterator gets raised in the main thread

Signed-off-by: Aanand Prasad <[email protected]>

Conflicts:
	compose/cli/multiplexer.py
Aanand Prasad 10 年之前
父節點
當前提交
80d90a745a
共有 2 個文件被更改,包括 28 次插入5 次删除
  1. 11 5
      compose/cli/multiplexer.py
  2. 17 0
      tests/unit/multiplexer_test.py

+ 11 - 5
compose/cli/multiplexer.py

@@ -26,7 +26,11 @@ class Multiplexer(object):
 
         while self._num_running > 0:
             try:
-                item = self.queue.get(timeout=0.1)
+                item, exception = self.queue.get(timeout=0.1)
+
+                if exception:
+                    raise exception
+
                 if item is STOP:
                     self._num_running -= 1
                 else:
@@ -42,7 +46,9 @@ class Multiplexer(object):
 
 
 def _enqueue_output(iterator, queue):
-    for item in iterator:
-        queue.put(item)
-
-    queue.put(STOP)
+    try:
+        for item in iterator:
+            queue.put((item, None))
+        queue.put((STOP, None))
+    except Exception as e:
+        queue.put((None, e))

+ 17 - 0
tests/unit/multiplexer_test.py

@@ -26,3 +26,20 @@ class MultiplexerTest(unittest.TestCase):
             [0, 1, 2, 3, 4, 5],
             sorted(list(mux.loop())),
         )
+
+    def test_exception(self):
+        class Problem(Exception):
+            pass
+
+        def problematic_iterator():
+            yield 0
+            yield 2
+            raise Problem(":(")
+
+        mux = Multiplexer([
+            problematic_iterator(),
+            (x for x in [1, 3, 5]),
+        ])
+
+        with self.assertRaises(Problem):
+            list(mux.loop())