Repository: dask/dask

Issue:
Define __await__ method for collections
What should this do?

```python
async def f():
    df = dask.datasets.timeseries()
    df = await df
```

One option, would be that it is shorthand for the following:

```python
async def f():
    df = dask.datasets.timeseries()
    df = df.persist()
    await wait(df)
```


Discussion:
That's the most reasonable behavior I see.

On Fri, Jun 7, 2019 at 12:21 PM Matthew Rocklin <notifications@github.com>
wrote:

> What should this do?
>
> async def f():
>     df = dask.datasets.timeseries()
>     df = await df
>
> One option, would be that it is shorthand for the following:
>
> async def f():
>     df = dask.datasets.timeseries()
>     df = df.persist()
>     await wait(df)
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/4900?email_source=notifications&email_token=AAKAOIXBYY7EQC6DU5CUBE3PZKKJJA5CNFSM4HVY43Q2YY3PNVWWK3TUL52HS4DFUVEXG43VMWVGG33NNVSW45C7NFSM4GYJUFBQ>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AAKAOIXE743D6U5PFAOVX7TPZKKJJANCNFSM4HVY43QQ>
> .
>


Resolving patch:
diff --git a/dask/base.py b/dask/base.py
--- a/dask/base.py
+++ b/dask/base.py
@@ -156,6 +156,23 @@ def compute(self, **kwargs):
         (result,) = compute(self, traverse=False, **kwargs)
         return result
 
+    def __await__(self):
+        try:
+            from distributed import wait, futures_of
+        except ImportError:
+            raise ImportError(
+                "Using async/await with dask requires the `distributed` package"
+            )
+        from tornado import gen
+
+        @gen.coroutine
+        def f():
+            if futures_of(self):
+                yield wait(self)
+            raise gen.Return(self)
+
+        return f().__await__()
+
 
 def compute_as_if_collection(cls, dsk, keys, scheduler=None, get=None, **kwargs):
     """Compute a graph as if it were of type cls.


Test patch:
diff --git a/dask/tests/py3_test_await.py b/dask/tests/py3_test_await.py
new file mode 100644
--- /dev/null
+++ b/dask/tests/py3_test_await.py
@@ -0,0 +1,24 @@
+# flake8: noqa  # while we support Python 2
+import asyncio
+import dask
+from distributed.utils_test import gen_cluster, inc
+from distributed import futures_of
+
+
+@gen_cluster(client=True)
+async def test_await(c, s, a, b):
+    x = dask.delayed(inc)(1)
+    x = await x.persist()
+    assert x.key in s.tasks
+    assert a.data or b.data
+    assert all(f.done() for f in futures_of(x))
+
+
+def test_local_scheduler():
+    async def f():
+        x = dask.delayed(inc)(1)
+        y = x + 1
+        z = await y.persist()
+        assert len(z.dask) == 1
+
+    asyncio.get_event_loop().run_until_complete(f())
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -680,7 +680,6 @@ def test_visualize_order():
 
 
 def test_use_cloudpickle_to_tokenize_functions_in__main__():
-    import sys
     from textwrap import dedent
 
     defn = dedent("""
@@ -955,3 +954,7 @@ def f(x):
     workers = {i.worker_id for i in prof.results}
 
     assert len(workers) == num_workers
+
+
+if sys.version_info >= (3, 5):
+    from dask.tests.py3_test_await import *  # noqa F401
diff --git a/dask/tests/test_utils.py b/dask/tests/test_utils.py
--- a/dask/tests/test_utils.py
+++ b/dask/tests/test_utils.py
@@ -423,6 +423,7 @@ def f(a, c):
 @pytest.mark.skipif(PY2, reason="Docstrings not as easy to manipulate in Py2")
 def test_derived_from_func():
     import builtins
+
     @derived_from(builtins)
     def sum():
         "extra docstring"