Simpler context manager implementation.

ajdavis · Jan 31, 2015 · e081117 · e081117
1 parent 7edcca2
commit e081117
Showing 1 changed file with 2 additions and 8 deletions.
diff --git a/toro/locks.py b/toro/locks.py
@@ -5,7 +5,7 @@
 import contextlib
 import collections
 
-from tornado import ioloop
+from tornado import gen, ioloop
 from tornado.concurrent import Future
 
 from . import _util
@@ -25,15 +25,9 @@ class _ContextManagerFuture(Future):
     """
     def __init__(self, wrapped, exit_callback):
         super(_ContextManagerFuture, self).__init__()
-        wrapped.add_done_callback(self._done_callback)
+        gen.chain_future(wrapped, self)
         self.exit_callback = exit_callback
 
-    def _done_callback(self, wrapped):
-        if wrapped.exception():
-            self.set_exception(wrapped.exception())
-        else:
-            self.set_result(wrapped.result())
-
     def result(self):
         if self.exception():
             raise self.exception()