From aaab17c3190b5a2978832063e4b3923cc7c2ebbc Mon Sep 17 00:00:00 2001 From: Hendrik Makait Date: Wed, 5 Oct 2022 15:13:55 +0200 Subject: [PATCH] Return default on KeyError at any level in get_metadata (#7109) --- distributed/scheduler.py | 6 +++--- distributed/tests/test_client.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/distributed/scheduler.py b/distributed/scheduler.py index f0e1225756..52c07d81b9 100644 --- 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 index aa8a001ac5..142c746e73 100644 --- 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)