Support PEP 585 dict syntax in dagster-k8s models

Author: vanika02

python_modules/libraries/dagster-k8s/dagster_k8s/models.py

@@ -39,12 +39,16 @@ def _get_openapi_list_element_type(attr_type: str) -> str:
 
 
 def _is_openapi_dict_type(attr_type: str) -> bool:
-    return attr_type.startswith("dict(")
+    return attr_type.startswith("dict(") or attr_type.startswith("dict[")
 
 
 def _get_openapi_dict_value_type(attr_type: str) -> str:
     # group(2) because value, not key
-    match = check.not_none(re.match(r"dict\(([^,]*), (.*)\)", attr_type))
+    match = check.not_none(
+        re.match(r"dict\(([^,]*), (.*)\)", attr_type)
+        or re.match(r"dict\[([^,]*),\s*(.*)\]", attr_type)
+    )
+
     return match.group(2)
 
 

python_modules/libraries/dagster-k8s/dagster_k8s_tests/unit_tests/test_models.py

@@ -150,3 +150,36 @@ def test_logger_calls_bounded():
 
     # Creating a model should not use the logger lock
     assert mock_set_level.call_count == 0
+
+
+def test_snake_case_dict_type_pep585_format():
+    with mock.patch.object(
+        kubernetes.client.V1CSIVolumeSource,
+        "openapi_types",
+        {
+            **kubernetes.client.V1CSIVolumeSource.openapi_types,
+            "volume_attributes": "dict[str, str]",
+        },
+    ):
+        volume_dict = {
+            "name": "my_volume",
+            "csi": {
+                "driver": "my_driver",
+                "volumeAttributes": {
+                    "fooKey": "fooVal"
+                },
+            },
+        }
+
+        assert k8s_snake_case_dict(
+            kubernetes.client.V1Volume,
+            volume_dict,
+        ) == {
+            "name": "my_volume",
+            "csi": {
+                "driver": "my_driver",
+                "volume_attributes": {
+                    "fooKey": "fooVal",
+                },
+            },
+        }