fix: support nested sequences in parallel yield for tasklets (#358)

This adds backwards compatibility with the legacy code.

Fixes #349.

Author: Chris Rossi

packages/google-cloud-ndb/google/cloud/ndb/tasklets.py

@@ -406,9 +406,15 @@ class _MultiFuture(Future):
 
     def __init__(self, dependencies):
         super(_MultiFuture, self).__init__()
-        self._dependencies = dependencies
-
+        futures = []
         for dependency in dependencies:
+            if isinstance(dependency, (list, tuple)):
+                dependency = _MultiFuture(dependency)
+            futures.append(dependency)
+
+        self._dependencies = futures
+
+        for dependency in futures:
             dependency.add_done_callback(self._dependency_done)
 
         if not dependencies:

packages/google-cloud-ndb/tests/system/test_crud.py

@@ -208,6 +208,50 @@ def get_two_entities():
     assert entity2.bar == "naan"
 
 
+@pytest.mark.usefixtures("client_context")
+def test_retrieve_entities_in_parallel_nested(ds_entity):
+    """Regression test for #357.
+
+    https://github.com/googleapis/python-ndb/issues/357
+    """
+    entity1_id = test_utils.system.unique_resource_id()
+    ds_entity(KIND, entity1_id, foo=42, bar="none")
+    entity2_id = test_utils.system.unique_resource_id()
+    ds_entity(KIND, entity2_id, foo=65, bar="naan")
+    entity3_id = test_utils.system.unique_resource_id()
+    ds_entity(KIND, entity3_id, foo=66, bar="route")
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+        bar = ndb.StringProperty()
+
+    key1 = ndb.Key(KIND, entity1_id)
+    key2 = ndb.Key(KIND, entity2_id)
+    key3 = ndb.Key(KIND, entity3_id)
+
+    @ndb.tasklet
+    def get_two_entities():
+        entity1, (entity2, entity3) = yield (
+            key1.get_async(),
+            [key2.get_async(), key3.get_async()],
+        )
+        raise ndb.Return(entity1, entity2, entity3)
+
+    entity1, entity2, entity3 = get_two_entities().result()
+
+    assert isinstance(entity1, SomeKind)
+    assert entity1.foo == 42
+    assert entity1.bar == "none"
+
+    assert isinstance(entity2, SomeKind)
+    assert entity2.foo == 65
+    assert entity2.bar == "naan"
+
+    assert isinstance(entity3, SomeKind)
+    assert entity3.foo == 66
+    assert entity3.bar == "route"
+
+
 @pytest.mark.usefixtures("client_context")
 def test_insert_entity(dispose_of, ds_client):
     class SomeKind(ndb.Model):

packages/google-cloud-ndb/tests/unit/test_tasklets.py

@@ -458,6 +458,15 @@ def test_no_dependencies():
         future = tasklets._MultiFuture(())
         assert future.result() == ()
 
+    @staticmethod
+    def test_nested():
+        dependencies = [tasklets.Future() for _ in range(3)]
+        future = tasklets._MultiFuture((dependencies[0], dependencies[1:]))
+        for i, dependency in enumerate(dependencies):
+            dependency.set_result(i)
+
+        assert future.result() == (0, (1, 2))
+
 
 class Test__get_return_value:
     @staticmethod