Implement tasklets.toplevel.

src/google/cloud/ndb/tasklets.py

@@ -559,5 +559,25 @@ def synctasklet_wrapper(*args, **kwargs):
     return synctasklet_wrapper
 
 
-def toplevel(*args, **kwargs):
-    raise NotImplementedError
+def toplevel(wrapped):
+    """A synctasklet decorator that flushes any pending work.
+
+    Use of this decorator is largely unnecessary, as you should be using
+    :meth:`~google.cloud.ndb.client.Client.context` which also flushes pending
+    work when exiting the context.
+
+    Args:
+        wrapped (Callable): The wrapped function."
+    """
+    synctasklet_wrapped = synctasklet(wrapped)
+
+    @functools.wraps(wrapped)
+    def toplevel_wrapper(*args, **kwargs):
+        context = context_module.get_context()
+        try:
+            with context.new().use():
+                return synctasklet_wrapped(*args, **kwargs)
+        finally:
+            _eventloop.run()
+
+    return toplevel_wrapper

tests/system/test_query.py

@@ -80,13 +80,13 @@ def test_fetch_lots_of_a_kind(dispose_of):
     class SomeKind(ndb.Model):
         foo = ndb.IntegerProperty()
 
-    @ndb.tasklet
+    @ndb.toplevel
     def make_entities():
         entities = [SomeKind(foo=i) for i in range(n_entities)]
         keys = yield [entity.put_async() for entity in entities]
         return keys
 
-    for key in make_entities().result():
+    for key in make_entities():
         dispose_of(key._key)
 
     query = SomeKind.query()
@@ -228,7 +228,7 @@ class SomeKind(ndb.Model):
         foo = ndb.IntegerProperty()
         bar = ndb.StringProperty()
 
-    @ndb.tasklet
+    @ndb.toplevel
     def make_entities():
         keys = yield (
             SomeKind(foo=1, bar="a").put_async(),
@@ -238,7 +238,7 @@ def make_entities():
         for key in keys:
             dispose_of(key._key)
 
-    make_entities().check_success()
+    make_entities()
     eventually(SomeKind.query().fetch, _length_equals(3))
 
     query = SomeKind.query(ndb.OR(SomeKind.foo == 1, SomeKind.bar == "c"))
@@ -290,7 +290,7 @@ class SomeKind(ndb.Model):
         foo = ndb.IntegerProperty()
         bar = ndb.StringProperty()
 
-    @ndb.tasklet
+    @ndb.toplevel
     def make_entities():
         keys = yield (
             SomeKind(foo=0, bar="a").put_async(),
@@ -301,7 +301,7 @@ def make_entities():
         for key in keys:
             dispose_of(key._key)
 
-    make_entities().check_success()
+    make_entities()
     query = SomeKind.query(ndb.OR(SomeKind.bar == "a", SomeKind.bar == "b"))
     query = query.order(SomeKind.foo)
     results = eventually(query.fetch, _length_equals(4))
@@ -353,7 +353,7 @@ class SomeKind(ndb.Model):
         foo = ndb.IntegerProperty()
         bar = ndb.StringProperty()
 
-    @ndb.tasklet
+    @ndb.toplevel
     def make_entities():
         keys = yield (
             SomeKind(foo=0, bar="a").put_async(),
@@ -366,7 +366,7 @@ def make_entities():
         for key in keys:
             dispose_of(key._key)
 
-    make_entities().check_success()
+    make_entities()
     eventually(SomeKind.query().fetch, _length_equals(6))
 
     query = SomeKind.query(ndb.OR(SomeKind.bar == "a", SomeKind.bar == "b"))
@@ -498,13 +498,13 @@ def test_fetch_page(dispose_of):
     class SomeKind(ndb.Model):
         foo = ndb.IntegerProperty()
 
-    @ndb.tasklet
+    @ndb.toplevel
     def make_entities():
         entities = [SomeKind(foo=i) for i in range(n_entities)]
         keys = yield [entity.put_async() for entity in entities]
         return keys
 
-    for key in make_entities().result():
+    for key in make_entities():
         dispose_of(key._key)
 
     query = SomeKind.query().order(SomeKind.foo)

tests/unit/test_tasklets.py

@@ -606,6 +606,18 @@ def generator_function(value):
     assert result == 11
 
 
+@pytest.mark.usefixtures("in_context")
 def test_toplevel():
-    with pytest.raises(NotImplementedError):
-        tasklets.toplevel()
+    @tasklets.toplevel
+    def generator_function(value):
+        future = tasklets.Future(value)
+        future.set_result(value)
+        x = yield future
+        return x + 3
+
+    idle = mock.Mock(__name__="idle", return_value=None)
+    _eventloop.add_idle(idle)
+
+    result = generator_function(8)
+    assert result == 11
+    idle.assert_called_once_with()