Return default on KeyError at any level in get_metadata (#7109)

dask · Oct 5, 2022 · aaab17c · aaab17c
1 parent 61353b7
commit aaab17c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
@@ -6939,10 +6939,10 @@ def set_metadata(self, keys: list[str], value: object = None) -> None:
 
     def get_metadata(self, keys: list[str], default=no_default):
         metadata = self.task_metadata
-        for key in keys[:-1]:
-            metadata = metadata[key]
         try:
-            return metadata[keys[-1]]
+            for key in keys:
+                metadata = metadata[key]
+            return metadata
         except KeyError:
             if default != no_default:
                 return default

diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
@@ -5583,10 +5583,29 @@ def fib(x):
 
 @gen_cluster(client=True)
 async def test_task_metadata(c, s, a, b):
+    with pytest.raises(KeyError):
+        await c.get_metadata("x")
+    with pytest.raises(KeyError):
+        await c.get_metadata(["x"])
+    result = await c.get_metadata("x", None)
+    assert result is None
+    result = await c.get_metadata(["x"], None)
+    assert result is None
+
+    with pytest.raises(KeyError):
+        await c.get_metadata(["x", "y"])
+    result = await c.get_metadata(["x", "y"], None)
+    assert result is None
+
     await c.set_metadata("x", 1)
     result = await c.get_metadata("x")
     assert result == 1
 
+    with pytest.raises(TypeError):
+        await c.get_metadata(["x", "y"])
+    with pytest.raises(TypeError):
+        await c.get_metadata(["x", "y"], None)
+
     future = c.submit(inc, 1)
     key = future.key
     await wait(future)