Bugfix Omegaconf plugin: Properly deal with NoneType values (#3056)

* Treat builtins.NoneType explicitly when parsing type descriptions Signed-off-by: Adrian Loy <[email protected]> * Modified tests to cover NoneType case Signed-off-by: Adrian Loy <[email protected]> --------- Signed-off-by: Adrian Loy <[email protected]>
flyteorg · Jan 23, 2025 · 06db981 · 06db981
1 parent 8a6bbd0
commit 06db981
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git 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
@@ -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]),
     ],