Issue #131: as_completed() and wait() now raises a TypeError if the list of

vstinner · vstinner · commit aa40ec5a0f61 · 2014-02-11T11:53:15.000+01:00
futures is not a list but a Future, Task or coroutine object
diff --git a/asyncio/tasks.py b/asyncio/tasks.py
@@ -358,6 +358,8 @@ def wait(fs, *, loop=None, timeout=None, return_when=ALL_COMPLETED):
     Note: This does not raise TimeoutError! Futures that aren't done
     when the timeout occurs are returned in the second set.
     """
+    if isinstance(fs, futures.Future) or iscoroutine(fs):
+        raise TypeError("expect a list of futures, not %s" % type(fs).__name__)
     if not fs:
         raise ValueError('Set of coroutines/Futures is empty.')
 
@@ -474,6 +476,8 @@ def as_completed(fs, *, loop=None, timeout=None):
 
     Note: The futures 'f' are not necessarily members of fs.
     """
+    if isinstance(fs, futures.Future) or iscoroutine(fs):
+        raise TypeError("expect a list of futures, not %s" % type(fs).__name__)
     loop = loop if loop is not None else events.get_event_loop()
     deadline = None if timeout is None else loop.time() + timeout
     todo = {async(f, loop=loop) for f in set(fs)}
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
@@ -7,6 +7,11 @@
 from asyncio import test_utils
 
 
+@asyncio.coroutine
+def coroutine_function():
+    pass
+
+
 class Dummy:
 
     def __repr__(self):
@@ -1338,6 +1343,27 @@ def test_gather_shield(self):
         child2.set_result(2)
         test_utils.run_briefly(self.loop)
 
+    def test_as_completed_invalid_args(self):
+        fut = asyncio.Future(loop=self.loop)
+
+        # as_completed() expects a list of futures, not a future instance
+        self.assertRaises(TypeError, self.loop.run_until_complete,
+            asyncio.as_completed(fut, loop=self.loop))
+        self.assertRaises(TypeError, self.loop.run_until_complete,
+            asyncio.as_completed(coroutine_function(), loop=self.loop))
+
+    def test_wait_invalid_args(self):
+        fut = asyncio.Future(loop=self.loop)
+
+        # wait() expects a list of futures, not a future instance
+        self.assertRaises(TypeError, self.loop.run_until_complete,
+            asyncio.wait(fut, loop=self.loop))
+        self.assertRaises(TypeError, self.loop.run_until_complete,
+            asyncio.wait(coroutine_function(), loop=self.loop))
+
+        # wait() expects at least a future
+        self.assertRaises(ValueError, self.loop.run_until_complete,
+            asyncio.wait([], loop=self.loop))
 
 class GatherTestsBase:
 
