Let Future objects be used as context managers.

2 files changed, 36 insertions, 41 deletions

diff --git a/dimension/client.py.in b/dimension/client.py.in
index 127b96c..2eaf3c3 100644
--- a/dimension/client.py.in
+++ b/dimension/client.py.in
@@ -157,28 +157,27 @@ def main():
     scene.adc_bailout = float(args.adc_bailout)
 
   # Ray-trace the scene
-  future = scene.ray_trace_async()
-  bar = None
-  if not args.quiet:
-    if scene.nthreads == 1:
-      render_message = "Rendering scene"
-    else:
-      render_message = "Rendering scene (using %d threads)" % scene.nthreads
-    bar = progress_bar_async(render_message, future)
-  if args.preview:
-    if have_preview:
-      preview.show_preview(canvas, future)
-    else:
-      print("Couldn't display preview window", file = sys.stderr)
-  bar.join()
-  future.join()
+  with scene.ray_trace_async() as future:
+    bar = None
+    if not args.quiet:
+      if scene.nthreads == 1:
+        render_message = "Rendering scene"
+      else:
+        render_message = "Rendering scene (using %d threads)" % scene.nthreads
+      bar = progress_bar_async(render_message, future)
+    if args.preview:
+      if have_preview:
+        preview.show_preview(canvas, future)
+      else:
+        print("Couldn't display preview window", file = sys.stderr)
+    if bar is not None:
+      bar.join()
 
   # Write the output file
   export_timer = Timer()
-  future = canvas.write_PNG_async(args.output)
-  if not args.quiet:
-    progress_bar("Writing %s" % args.output, future)
-  future.join()
+  with canvas.write_PNG_async(args.output) as future:
+    if not args.quiet:
+      progress_bar("Writing %s" % args.output, future)
   export_timer.stop()
 
   # Print execution times
@@ -233,30 +232,18 @@ def working_directory(newwd):
 
 def progress_bar(str, future):
   """Display a progress bar while a Future completes."""
-  try:
-    print(str, end = " ")
-    sys.stdout.flush()
-
-    term_width = terminal_width()
-    width = term_width - (len(str) + 1)%term_width
-    for i in range(width):
-      future.wait((i + 1)/width)
-      print(".", end = "")
-      sys.stdout.flush()
-
-    print()
-    sys.stdout.flush()
-  except KeyboardInterrupt:
-    print()
+  print(str, end = " ")
+  sys.stdout.flush()
+
+  term_width = terminal_width()
+  width = term_width - (len(str) + 1)%term_width
+  for i in range(width):
+    future.wait((i + 1)/width)
+    print(".", end = "")
     sys.stdout.flush()
 
-    future.cancel()
-    try:
-      future.join()
-    except RuntimeError:
-      # Swallow the failure exception
-      pass
-    raise
+  print()
+  sys.stdout.flush()
 
 def progress_bar_async(str, future):
   thread = threading.Thread(target = progress_bar, args = (str, future))
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.")