Add afar.get, which automatically gathers data locally

README.md

@@ -21,20 +21,27 @@ Outside the context, `result` is a [Dask Future](https://docs.dask.org/en/latest
 
 By default, only the last assignment is saved.  One can specify which variables to save:
 ```python
-with afar.run("a", "b"), remotely:
-    a = 1
-    b = a + 1
+with afar.run("one", "two"), remotely:
+    one = 1
+    two = one + 1
 ```
-`a` and `b` are now both Futures.  They can be used directly in other `afar.run` contexts:
+`one` and `two` are now both Futures.  They can be used directly in other `afar.run` contexts:
 ```python
 with afar.run as data, remotely:
-    c = a + b
+    three = one + two
 
-assert c.result() == 3
-assert data["c"].result() == 3
+assert three.result() == 3
+assert data["three"].result() == 3
 ```
 `data` is a dictionary of variable names to Futures.  It may be necessary at times to get the data from here.
 
+If you want to automatically gather the data locally (to avoid calling `.result()`), use `afar.get` instead of `afar.run`:
+```python
+with afar.get, remotely:
+    five = two + three
+assert five == 5
+```
+
 ### Is this a good idea?
 
 I don't know!

afar/__init__.py

@@ -1,4 +1,4 @@
-from .core import run, remotely, locally  # noqa
+from .core import run, get, remotely, locally  # noqa
 
 from ._version import get_versions
 

afar/core.py

@@ -41,6 +41,8 @@ def __call__(self, **submit_kwargs):
 
 
 class Run:
+    _gather_data = False
+
     def __init__(self, *names):
         self.names = names
         self._results = None
@@ -51,7 +53,7 @@ def __init__(self, *names):
         self._with_lineno = None
 
     def __call__(self, *names):
-        return Run(*names)
+        return type(self)(*names)
 
     def __enter__(self):
         self._frame = inspect.currentframe().f_back
@@ -149,8 +151,16 @@ def _exit(self, exc_type, exc_value, exc_traceback):
             remote_dict = client.submit(
                 run_on_worker, self._scoped, names, futures, **submit_kwargs
             )
-            for name in names:
-                self._results[name] = client.submit(getitem, remote_dict, name, **submit_kwargs)
+            if self._gather_data:
+                futures_to_name = {
+                    client.submit(getitem, remote_dict, name, **submit_kwargs): name
+                    for name in names
+                }
+                for future, result in distributed.as_completed(futures_to_name, with_results=True):
+                    self._results[futures_to_name[future]] = result
+            else:
+                for name in names:
+                    self._results[name] = client.submit(getitem, remote_dict, name, **submit_kwargs)
         else:
             # Run locally.  This is handy for testing and debugging.
             results = self._scoped()
@@ -163,10 +173,17 @@ def _exit(self, exc_type, exc_value, exc_traceback):
         return True
 
 
+class Get(Run):
+    """Unlike ``run``, ``get`` automatically gathers the data locally"""
+
+    _gather_data = True
+
+
 def run_on_worker(sfunc, names, futures):
     sfunc = sfunc.bind(futures)
     results = sfunc()
     return {key: results[key] for key in names}
 
 
 run = Run()
+get = Get()

afar/tests/test_remotely.py

@@ -1,13 +1,23 @@
-import afar
-from operator import add
-from dask.distributed import Client
+import subprocess
+import sys
 
 # TODO: better testing infrastructure
 if __name__ == "__main__":
+    import afar
+    from operator import add
+    from dask.distributed import Client
+
     client = Client()
     two = client.submit(add, 1, 1)
 
-    with afar.run, remotely:
+    with afar.run as results, afar.remotely:
         three = two + 1
-
     assert three.result() == 3
+
+    with afar.get, afar.remotely(priority=1):
+        five = two + three
+    assert five == 5
+
+
+def test_runme():
+    assert subprocess.check_call([sys.executable, __file__]) == 0