From 06db981a5b8e687adb360131f27793f4069edd04 Mon Sep 17 00:00:00 2001 From: adrianloy Date: Thu, 23 Jan 2025 02:13:29 +0100 Subject: [PATCH] Bugfix Omegaconf plugin: Properly deal with NoneType values (#3056) * Treat builtins.NoneType explicitly when parsing type descriptions Signed-off-by: Adrian Loy * Modified tests to cover NoneType case Signed-off-by: Adrian Loy --------- Signed-off-by: Adrian Loy --- .../flytekitplugins/omegaconf/dictconfig_transformer.py | 2 ++ .../tests/test_dictconfig_transformer.py | 7 ++++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/plugins/flytekit-omegaconf/flytekitplugins/omegaconf/dictconfig_transformer.py b/plugins/flytekit-omegaconf/flytekitplugins/omegaconf/dictconfig_transformer.py index 0f2b8c63cc..52827f9baa 100644 --- a/plugins/flytekit-omegaconf/flytekitplugins/omegaconf/dictconfig_transformer.py +++ b/plugins/flytekit-omegaconf/flytekitplugins/omegaconf/dictconfig_transformer.py @@ -143,6 +143,8 @@ def parse_type_description(type_desc: str) -> Type: return origin[sub_types[0]] return origin[tuple(sub_types)] else: + if type_desc == "builtins.NoneType": + return NoneType module_name, class_name = type_desc.rsplit(".", 1) return importlib.import_module(module_name).__getattribute__(class_name) diff --git a/plugins/flytekit-omegaconf/tests/test_dictconfig_transformer.py b/plugins/flytekit-omegaconf/tests/test_dictconfig_transformer.py index b4d9115fa9..6ab531cb22 100644 --- a/plugins/flytekit-omegaconf/tests/test_dictconfig_transformer.py +++ b/plugins/flytekit-omegaconf/tests/test_dictconfig_transformer.py @@ -5,7 +5,7 @@ check_if_valid_dictconfig, extract_type_and_value_maps, is_flattenable, - parse_type_description, + parse_type_description, NoneType, ) from omegaconf import DictConfig, OmegaConf @@ -77,11 +77,11 @@ def test_is_flattenable(config: DictConfig, should_flatten: bool, monkeypatch: p def test_extract_type_and_value_maps_simple() -> None: """Test extraction of type and value maps from a simple DictConfig.""" ctx = FlyteContext.current_context() - config: DictConfig = OmegaConf.create({"key1": "value1", "key2": 123, "key3": True}) + config: DictConfig = OmegaConf.create({"key1": "value1", "key2": 123, "key3": True, "key4": None}) type_map, value_map = extract_type_and_value_maps(ctx, config) - expected_type_map = {"key1": "builtins.str", "key2": "builtins.int", "key3": "builtins.bool"} + expected_type_map = {"key1": "builtins.str", "key2": "builtins.int", "key3": "builtins.bool", "key4": "builtins.NoneType"} assert type_map == expected_type_map assert "key1" in value_map @@ -93,6 +93,7 @@ def test_extract_type_and_value_maps_simple() -> None: "type_desc, expected_type", [ ("builtins.int", int), + ("builtins.NoneType", NoneType), ("typing.List[builtins.int]", t.List[int]), ("typing.Optional[builtins.int]", t.Optional[int]), ],