Make Greenlets context managers to handle their lifetime.

Fixes #1324

docs/api/gevent.greenlet.rst

@@ -17,20 +17,34 @@ Starting Greenlets
 To start a new greenlet, pass the target function and its arguments to
 :class:`Greenlet` constructor and call :meth:`Greenlet.start`:
 
->>> g = Greenlet(myfunction, 'arg1', 'arg2', kwarg1=1)
->>> g.start()
+   >>> from gevent import Greenlet
+   >>> def myfunction(arg1, arg2, kwarg1=None):
+   ...    pass
+   >>> g = Greenlet(myfunction, 'arg1', 'arg2', kwarg1=1)
+   >>> g.start()
 
 or use classmethod :meth:`Greenlet.spawn` which is a shortcut that
 does the same:
 
->>> g = Greenlet.spawn(myfunction, 'arg1', 'arg2', kwarg1=1)
+   >>> g = Greenlet.spawn(myfunction, 'arg1', 'arg2', kwarg1=1)
 
 There are also various spawn helpers in :mod:`gevent`, including:
 
 - :func:`gevent.spawn`
 - :func:`gevent.spawn_later`
 - :func:`gevent.spawn_raw`
 
+Waiting For Greenlets
+=====================
+
+You can wait for a greenlet to finish with its :meth:`Greenlet.join`
+method. There are helper functions to join multiple greenlets or
+heterogenous collections of objects:
+
+- :func:`gevent.joinall`
+- :func:`gevent.wait`
+- :func:`gevent.iwait`
+
 Stopping Greenlets
 ==================
 
@@ -41,6 +55,48 @@ circumstances (if you might have a :class:`raw greenlet <greenlet.greenlet>`):
 - :func:`gevent.kill`
 - :func:`gevent.killall`
 
+Context Managers
+================
+
+.. versionadded:: 21.1.0
+
+
+Greenlets also function as context managers, so you can combine
+spawning and waiting for a greenlet to finish in a single line:
+
+.. doctest::
+
+    >>> def in_greenlet():
+    ...     print("In the greenlet")
+    ...     return 42
+    >>> with Greenlet.spawn(in_greenlet) as g:
+    ...     print("In the with suite")
+    In the with suite
+    In the greenlet
+    >>> g.get(block=False)
+    42
+
+Normally, the greenlet is joined to wait for it to finish, but if the
+body of the suite raises an exception, the greenlet is killed with
+that exception.
+
+.. doctest::
+
+   >>> import gevent
+   >>> try:
+   ...    with Greenlet.spawn(gevent.sleep, 0.1) as g:
+   ...       raise Exception("From with body")
+   ... except Exception:
+   ...     pass
+   >>> g.dead
+   True
+   >>> g.successful()
+   False
+   >>> g.get(block=False)
+   Traceback (most recent call last):
+   ...
+   Exception: From with body
+
 .. _subclassing-greenlet:
 
 Subclassing Greenlet

docs/changes/1324.feature

@@ -0,0 +1,7 @@
+Make :class:`gevent.Greenlet` objects function as context managers.
+When the ``with`` suite finishes, execution doesn't continue until the
+greenlet is finished. This can be a simpler alternative to a
+:class:`gevent.pool.Group` when the lifetime of greenlets can be
+lexically scoped.
+
+Suggested by André Caron.

docs/intro.rst

@@ -14,7 +14,7 @@ The following example shows how to run tasks concurrently.
     >>> from gevent import socket
     >>> urls = ['www.google.com', 'www.example.com', 'www.python.org']
     >>> jobs = [gevent.spawn(socket.gethostbyname, url) for url in urls]
-    >>> gevent.joinall(jobs, timeout=2)
+    >>> _ = gevent.joinall(jobs, timeout=2)
     >>> [job.value for job in jobs]
     ['74.125.79.106', '208.77.188.166', '82.94.164.162']
 
@@ -44,7 +44,7 @@ counterparts. That way even the modules that are unaware of gevent can benefit f
 in a multi-greenlet environment.
 
     >>> from gevent import monkey; monkey.patch_socket()
-    >>> import urllib2 # it's usable from multiple greenlets now
+    >>> import requests # it's usable from multiple greenlets now
 
 See :doc:`examples/concurrent_download`.
 
@@ -170,7 +170,7 @@ If there is an error during execution it won't escape the greenlet's
 boundaries. An unhandled error results in a stacktrace being printed,
 annotated by the failed function's signature and arguments:
 
-    >>> gevent.spawn(lambda : 1/0)
+    >>> glet = gevent.spawn(lambda : 1/0); glet.join()
     >>> gevent.sleep(1)
     Traceback (most recent call last):
      ...
@@ -195,6 +195,7 @@ Greenlets can be killed synchronously from another greenlet. Killing
 will resume the sleeping greenlet, but instead of continuing
 execution, a :exc:`GreenletExit` will be raised.
 
+    >>> from gevent import Greenlet
     >>> g = Greenlet(gevent.sleep, 4)
     >>> g.start()
     >>> g.kill()
@@ -225,10 +226,10 @@ catch it), thus it's a good idea always to pass a timeout to
 :meth:`kill <gevent.Greenlet.kill>` (otherwise, the greenlet doing the
 killing will remain blocked forever).
 
-.. tip:: The exact timing at which an exception is raised within a
-		  target greenlet as the result of :meth:`kill
-		  <gevent.Greenlet.kill>` is not defined. See that function's
-		  documentation for more details.
+.. tip::
+   The exact timing at which an exception is raised within a target
+   greenlet as the result of :meth:`kill <gevent.Greenlet.kill>` is
+   not defined. See that function's documentation for more details.
 
 .. caution::
    Use care when killing greenlets, especially arbitrary
@@ -249,6 +250,22 @@ killing will remain blocked forever).
    <http://docs.oracle.com/javase/8/docs/technotes/guides/concurrency/threadPrimitiveDeprecation.html>`_
    describes a similar situation for threads.
 
+
+Greenlets also function as context managers, so you can combine
+spawning and waiting for a greenlet to finish in a single line:
+
+.. doctest::
+
+    >>> def in_greenlet():
+    ...     print("In the greenlet")
+    ...     return 42
+    >>> with Greenlet.spawn(in_greenlet) as g:
+    ...     print("In the with suite")
+    In the with suite
+    In the greenlet
+    >>> g.get(block=False)
+    42
+
 Timeouts
 ========
 

src/gevent/_gevent_cgreenlet.pxd

@@ -128,6 +128,7 @@ cdef class Greenlet(greenlet):
 
     cpdef bint has_links(self)
     cpdef join(self, timeout=*)
+    cpdef kill(self, exception=*, block=*, timeout=*)
     cpdef bint ready(self)
     cpdef bint successful(self)
     cpdef rawlink(self, object callback)

src/gevent/_waiter.py

@@ -38,6 +38,8 @@ class Waiter(object):
     The :meth:`switch` and :meth:`throw` methods must only be called from the :class:`Hub` greenlet.
     The :meth:`get` method must be called from a greenlet other than :class:`Hub`.
 
+        >>> from gevent.hub import Waiter
+        >>> from gevent import get_hub
         >>> result = Waiter()
         >>> timer = get_hub().loop.timer(0.1)
         >>> timer.start(result.switch, 'hello from Waiter')
@@ -48,6 +50,7 @@ class Waiter(object):
     If switch is called before the greenlet gets a chance to call :meth:`get` then
     :class:`Waiter` stores the value.
 
+        >>> from gevent.time import sleep
         >>> result = Waiter()
         >>> timer = get_hub().loop.timer(0.1)
         >>> timer.start(result.switch, 'hi from Waiter')

src/gevent/event.py

@@ -175,6 +175,7 @@ class AsyncResult(AbstractLinkable): # pylint:disable=undefined-variable
     To pass a value call :meth:`set`. Calls to :meth:`get` (those that are currently blocking as well as
     those made in the future) will return the value:
 
+        >>> from gevent.event import AsyncResult
         >>> result = AsyncResult()
         >>> result.set(100)
         >>> result.get()

src/gevent/greenlet.py

@@ -202,6 +202,13 @@ def __init__(self, run=None, *args, **kwargs):
         .. versionchanged:: 1.5
            Greenlet objects are now more careful to verify that their ``parent`` is really
            a gevent hub, raising a ``TypeError`` earlier instead of an ``AttributeError`` later.
+
+        .. versionchanged:: NEXT
+           Greenlet objects now function as context managers. Exiting the ``with`` suite
+           ensures that the greenlet has completed by :meth:`joining <join>`
+           the greenlet (blocking, with
+           no timeout). If the body of the suite raises an exception, the greenlet is
+           :meth:`killed <kill>` with the default arguments and not joined in that case.
         """
         # The attributes are documented in the .rst file
 
@@ -477,6 +484,8 @@ def __handle_death_before_start(self, args):
                 args = (GreenletExit, GreenletExit(), None)
             if not issubclass(args[0], BaseException):
                 # Random non-type, non-exception arguments.
+                print("RANDOM CRAP", args)
+                import traceback; traceback.print_stack()
                 args = (BaseException, BaseException(args), None)
             assert issubclass(args[0], BaseException)
             self.__report_error(args)
@@ -707,7 +716,11 @@ def _maybe_kill_before_start(self, exception):
         self.__free()
         dead = self.dead
         if dead:
-            self.__handle_death_before_start((exception,))
+            if isinstance(exception, tuple) and len(exception) == 3:
+                args = exception
+            else:
+                args = (exception,)
+            self.__handle_death_before_start(args)
         return dead
 
     def kill(self, exception=GreenletExit, block=True, timeout=None):
@@ -756,8 +769,14 @@ def kill(self, exception=GreenletExit, block=True, timeout=None):
             If this greenlet had never been switched to, killing it will
             prevent it from *ever* being switched to. Links (:meth:`rawlink`)
             will still be executed, though.
+        .. versionchanged:: NEXT
+            If this greenlet is :meth:`ready`, immediately return instead of
+            requiring a trip around the event loop.
         """
         if not self._maybe_kill_before_start(exception):
+            if self.ready():
+                return
+
             waiter = Waiter() if block else None # pylint:disable=undefined-variable
             hub = get_my_hub(self) # pylint:disable=undefined-variable
             hub.loop.run_callback(_kill, self, exception, waiter)
@@ -837,6 +856,18 @@ def join(self, timeout=None):
             self.unlink(switch)
             raise
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, t, v, tb):
+        if t is None:
+            try:
+                self.join()
+            finally:
+                self.kill()
+        else:
+            self.kill((t, v, tb))
+
     def __report_result(self, result):
         self._exc_info = (None, None, None)
         self.value = result
@@ -1012,7 +1043,10 @@ def close(self):
 # and its first argument is the Greenlet. So we can be sure about the types.
 def _kill(glet, exception, waiter):
     try:
-        glet.throw(exception)
+        if isinstance(exception, tuple) and len(exception) == 3:
+            glet.throw(*exception)
+        else:
+            glet.throw(exception)
     except: # pylint:disable=bare-except, undefined-variable
         # XXX do we need this here?
         get_my_hub(glet).handle_error(glet, *sys_exc_info())

src/gevent/local.py

@@ -11,6 +11,8 @@
 If you have data that you want to be local to a greenlet, simply create
 a greenlet-local object and use its attributes:
 
+  >>> import gevent
+  >>> from gevent.local import local
   >>> mydata = local()
   >>> mydata.number = 42
   >>> mydata.number

src/gevent/queue.py

@@ -13,6 +13,7 @@
 :meth:`get <Queue.get>` returns ``StopIteration`` (specifically that
 class, not an instance or subclass).
 
+    >>> import gevent.queue
     >>> queue = gevent.queue.Queue()
     >>> queue.put(1)
     >>> queue.put(2)

src/gevent/tests/test__greenlet.py

@@ -41,6 +41,27 @@
 class ExpectedError(greentest.ExpectedException):
     pass
 
+class ExpectedJoinError(ExpectedError):
+    pass
+
+class SuiteExpectedException(ExpectedError):
+    pass
+
+class GreenletRaisesJoin(gevent.Greenlet):
+    killed = False
+    joined = False
+    raise_on_join = True
+
+    def join(self, timeout=None):
+        self.joined += 1
+        if self.raise_on_join:
+            raise ExpectedJoinError
+        else:
+            return gevent.Greenlet.join(self, timeout)
+
+    def kill(self, *args, **kwargs): # pylint:disable=signature-differs
+        self.killed += 1
+        return gevent.Greenlet.kill(self, *args, **kwargs)
 
 class TestLink(greentest.TestCase):
 
@@ -879,6 +900,68 @@ def test_killall_raw(self):
         g = gevent.spawn_raw(lambda: 1)
         gevent.killall([g])
 
+
+class TestContextManager(greentest.TestCase):
+
+    def test_simple(self):
+        with gevent.spawn(gevent.sleep, timing.SMALL_TICK) as g:
+            self.assert_greenlet_spawned(g)
+        # It is completed after the suite
+        self.assert_greenlet_finished(g)
+
+    def test_wait_in_suite(self):
+        with gevent.spawn(self._raise_exception) as g:
+            with self.assertRaises(greentest.ExpectedException):
+                g.get()
+        self.assert_greenlet_finished(g)
+
+    @staticmethod
+    def _raise_exception():
+        raise greentest.ExpectedException
+
+    def test_greenlet_raises(self):
+        with gevent.spawn(self._raise_exception) as g:
+            pass
+
+        self.assert_greenlet_finished(g)
+        with self.assertRaises(greentest.ExpectedException):
+            g.get()
+
+    def test_join_raises(self):
+        suite_ran = 0
+        with self.assertRaises(ExpectedJoinError):
+            with GreenletRaisesJoin.spawn(gevent.sleep, timing.SMALL_TICK) as g:
+                self.assert_greenlet_spawned(g)
+                suite_ran = 1
+
+        self.assertTrue(suite_ran)
+        self.assert_greenlet_finished(g)
+        self.assertTrue(g.killed)
+
+    def test_suite_body_raises(self, delay=None):
+        greenlet_sleep = timing.SMALL_TICK if not delay else timing.LARGE_TICK
+        with self.assertRaises(SuiteExpectedException):
+            with GreenletRaisesJoin.spawn(gevent.sleep, greenlet_sleep) as g:
+                self.assert_greenlet_spawned(g)
+                if delay:
+                    g.raise_on_join = False
+                    gevent.sleep(delay)
+                raise SuiteExpectedException
+
+        self.assert_greenlet_finished(g)
+        self.assertTrue(g.killed)
+        if delay:
+            self.assertTrue(g.joined)
+        else:
+            self.assertFalse(g.joined)
+        self.assertFalse(g.successful())
+
+        with self.assertRaises(SuiteExpectedException):
+            g.get()
+
+    def test_suite_body_raises_with_delay(self):
+        self.test_suite_body_raises(delay=timing.SMALL_TICK)
+
 class TestStart(greentest.TestCase):
 
     def test_start(self):