Provide a context manager API to prevent memory leaks in gevent.iwait()

src/gevent/__hub_primitives.pxd

@@ -58,6 +58,10 @@ cdef class _WaitIterator:
     cdef _begin(self)
     cdef _cleanup(self)
 
+    cpdef __enter__(self)
+    cpdef __exit__(self, typ, value, tb)
+
+
 cpdef iwait_on_objects(objects, timeout=*, count=*)
 cpdef wait_on_objects(objects=*, timeout=*, count=*)
 

src/gevent/_hub_primitives.py

@@ -168,6 +168,12 @@ class _WaitIterator(object):
                 except: # pylint:disable=bare-except
                     traceback.print_exc()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, typ, value, tb):
+        self._cleanup()
+
 
 def iwait_on_objects(objects, timeout=None, count=None):
     """

src/greentest/test__iwait.py

@@ -16,5 +16,26 @@ class Testiwait(greentest.TestCase):
         ready = next(gevent.iwait((sem1, sem2)))
         self.assertEqual(sem1, ready)
 
+    def test_iwait_partial(self):
+        # Test that the iwait context manager allows the iterator to be
+        # consumed partially without a memory leak.
+
+        sem = Semaphore()
+        let = gevent.spawn(sem.release)
+        with gevent.iwait((sem,), timeout=0.01) as iterator:
+            self.assertEqual(sem, next(iterator))
+        let.get()
+
+    def test_iwait_nogarbage(self):
+        sem1 = Semaphore()
+        sem2 = Semaphore()
+        let = gevent.spawn(sem1.release)
+        with gevent.iwait((sem1, sem2)) as iterator:
+            self.assertEqual(sem1, next(iterator))
+            assert len(sem2._links) == 1
+        assert sem2._links is None or len(sem2._links) == 0
+        let.get()
+
+
 if __name__ == '__main__':
     greentest.main()