fix 'count' in wait/iwait/joinall

in joinall, if 'raise_error=False' and 'count is not None',
'count' is ignored. This commit passes this argument.
This commit also adds support for 'count' in 'iwait', which
simplifies the counting in joinall and wait.
In addition, to make the behavior of iwait and wait similar
(so that the behavior of joinall will not change according to the
value of 'raise_error'), handling objects=None was also added to 'iwait'

gevent/greenlet.py

@@ -410,15 +410,11 @@ def _kill(greenlet, exception, waiter):
 
 def joinall(greenlets, timeout=None, raise_error=False, count=None):
     if not raise_error:
-        wait(greenlets, timeout=timeout)
+        wait(greenlets, timeout=timeout, count=count)
     else:
-        for obj in iwait(greenlets, timeout=timeout):
+        for obj in iwait(greenlets, timeout=timeout, count=count):
             if getattr(obj, 'exception', None) is not None:
                 raise obj.exception
-            if count is not None:
-                count -= 1
-                if count <= 0:
-                    break
 
 
 def _killall3(greenlets, exception, waiter):

gevent/hub.py

@@ -593,19 +593,25 @@ class Waiter(object):
     # and unwraps it in wait() thus checking that switch() was indeed called
 
 
-def iwait(objects, timeout=None):
+def iwait(objects, timeout=None, count=None):
     """Yield objects as they are ready, until all are ready or timeout expired.
 
     *objects* must be iterable yielding instance implementing wait protocol (rawlink() and unlink()).
     """
     # QQQ would be nice to support iterable here that can be generated slowly (why?)
+    if objects is None:
+        yield get_hub().join(timeout=timeout)
+        return
     waiter = Waiter()
     switch = waiter.switch
     if timeout is not None:
         timer = get_hub().loop.timer(timeout, priority=-1)
         timer.start(waiter.switch, _NONE)
     try:
+        if count is None:
             count = len(objects)
+        else:
+            count = min(count, len(objects))
         for obj in objects:
             obj.rawlink(switch)
         for _ in xrange(count):
@@ -654,15 +660,7 @@ def wait(objects=None, timeout=None, count=None):
     """
     if objects is None:
         return get_hub().join(timeout=timeout)
-    result = []
-    if count is None:
-        return list(iwait(objects, timeout))
-    for obj in iwait(objects=objects, timeout=timeout):
-        result.append(obj)
-        count -= 1
-        if count <= 0:
-            break
-    return result
+    return list(iwait(objects, timeout, count))
 
 
 class linkproxy(object):