From 3e40d9c7b75dec008beb6b4ecf71b60e8bae088d Mon Sep 17 00:00:00 2001
From: Tavian Barnes <tavianator@gmail.com>
Date: Fri, 4 Nov 2011 19:02:27 -0400
Subject: Let Future objects be used as context managers.

---
 libdimension-python/wrapper.pyx | 8 ++++++++
 1 file changed, 8 insertions(+)

(limited to 'libdimension-python')

diff --git a/libdimension-python/wrapper.pyx b/libdimension-python/wrapper.pyx
index 1afb0ea..5f2b755 100644
--- a/libdimension-python/wrapper.pyx
+++ b/libdimension-python/wrapper.pyx
@@ -91,6 +91,14 @@ cdef class Future:
     with nogil:
       dmnsn_future_wait(self._future, progress)
 
+  # Let Futures be used as context managers
+  def __enter__(self):
+    return self
+  def __exit__(self, exc_type, exc_value, traceback):
+    if self._future != NULL:
+      self.join()
+    return False
+
   def _assert_unfinished(self):
     if self._future == NULL:
       raise RuntimeError("background task finished.")
-- 
cgit v1.2.3