Minor refactoring of PyDABs autogen code

python/codegen/codegen/generated_dataclass.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import TYPE_CHECKING, Optional
 
 if TYPE_CHECKING:
@@ -65,14 +65,6 @@ class GeneratedField:
     The type of the field in TypedDict, e.g., GeneratedType(name="TaskParam", ...)
     """
 
-    create_func_type_name: GeneratedType
-    """
-    Type type of the field in static "create" function, e.g., GeneratedType(name="TaskParam", ...)
-
-    It can be different from param_type_name because lists are made optional in "create" function
-    to avoid problems with mutable default arguments.
-    """
-
     description: Optional[str]
     """
     The description of the field to be included into a docstring.
@@ -83,14 +75,6 @@ class GeneratedField:
     The default value of the field, e.g., "None"
     """
 
-    create_func_default: Optional[str]
-    """
-    The default value of the field in "create" function.
-
-    It can be different from default because lists are made optional in "create" function
-    to avoid problems with mutable default arguments.
-    """
-
     default_factory: Optional[str]
     """
     Factory method for creating a default value, used for lists and dicts.
@@ -151,66 +135,27 @@ def generate_field(
     field_type = variable_or_type(field_type, is_required=is_required)
     param_type = variable_or_type(param_type, is_required=is_required)
 
+    # Base is the optional-scalar shape (default None). Collections instead use an
+    # empty-container factory, and required fields carry no default.
+    field = GeneratedField(
+        field_name=field_name,
+        type_name=field_type,
+        param_type_name=param_type,
+        description=prop.description,
+        default="None",
+        default_factory=None,
+        experimental=is_experimental_stage(prop.stage),
+        deprecated=prop.deprecated or False,
+    )
+
     if field_type.name == "VariableOrDict":
-        return GeneratedField(
-            field_name=field_name,
-            type_name=field_type,
-            param_type_name=param_type,
-            create_func_type_name=optional_type(param_type),
-            description=prop.description,
-            default=None,
-            default_factory="dict",
-            create_func_default="None",
-            experimental=is_experimental_stage(prop.stage),
-            deprecated=prop.deprecated or False,
-        )
+        return replace(field, default=None, default_factory="dict")
     elif field_type.name == "VariableOrList":
-        return GeneratedField(
-            field_name=field_name,
-            type_name=field_type,
-            param_type_name=param_type,
-            create_func_type_name=optional_type(param_type),
-            description=prop.description,
-            default=None,
-            default_factory="list",
-            create_func_default="None",
-            experimental=is_experimental_stage(prop.stage),
-            deprecated=prop.deprecated or False,
-        )
+        return replace(field, default=None, default_factory="list")
     elif is_required:
-        return GeneratedField(
-            field_name=field_name,
-            type_name=field_type,
-            param_type_name=param_type,
-            create_func_type_name=param_type,
-            description=prop.description,
-            default=None,
-            default_factory=None,
-            create_func_default=None,
-            experimental=is_experimental_stage(prop.stage),
-            deprecated=prop.deprecated or False,
-        )
-    else:
-        return GeneratedField(
-            field_name=field_name,
-            type_name=field_type,
-            param_type_name=param_type,
-            create_func_type_name=param_type,
-            description=prop.description,
-            default="None",
-            default_factory=None,
-            create_func_default="None",
-            experimental=is_experimental_stage(prop.stage),
-            deprecated=prop.deprecated or False,
-        )
+        return replace(field, default=None)
 
-
-def optional_type(generated: GeneratedType) -> GeneratedType:
-    return GeneratedType(
-        name="Optional",
-        package="typing",
-        parameters=[generated],
-    )
+    return field
 
 
 def str_type() -> GeneratedType:
@@ -281,7 +226,12 @@ def generate_type(namespace: str, ref: str, is_param: bool) -> GeneratedType:
             parameters=[element_type],
         )
 
-    if ref == "#/$defs/map/string":
+    if ref.startswith("#/$defs/map/"):
+        # Only dict[str, str] is modelled today; anything else (e.g. map/int,
+        # map/<object>) fails loudly instead.
+        if ref != "#/$defs/map/string":
+            raise ValueError(f"Unsupported map ref: {ref}")
+
         return dict_type()
 
     class_name = packages.get_class_name(ref)

python/codegen/codegen/generated_test_cases.py

@@ -1,8 +1,8 @@
 """
-Generates the per-resource TestCase data driving databricks_tests/core/test_resources.py.
+Generates the per-resource ResourceTestCase data driving databricks_tests/core/test_resources.py.
 
 For every wired resource a file _generated/<plural>.py is written (rendered from
-test_case.py.tmpl) exposing _test_case() -> (TestCase, _ResourceType). The generated
+test_case.py.tmpl) exposing _test_case() -> (ResourceTestCase, _ResourceType). The generated
 _generated/__init__.py collects them into `test_cases`, which test_resources.py imports
 and parametrizes its per-resource tests off.
 
@@ -160,8 +160,7 @@ def _synth_ref(
 
     schema = schemas[name]
     class_name = packages.get_class_name(ref)
-    module = packages.get_package(namespace, ref)
-    assert module
+    module = _module_of(namespace, ref)
 
     if schema.type == openapi.SchemaType.STRING:
         value = schema.enum[0]
@@ -208,7 +207,7 @@ def _synth_object(
                 continue
             if (
                 prop.deprecated
-                or _STAGE_RANK[prop.stage]
+                or _STAGE_RANK.get(prop.stage, 0)
                 > _STAGE_RANK[openapi.LaunchStage.PUBLIC_PREVIEW]
             ):
                 continue

python/codegen/codegen/main.py

@@ -53,7 +53,7 @@ def main(output: str):
     # decorators, and the core package __init__).
     generated_wiring.write_wiring(output)
 
-    # Generate the per-resource TestCase data driving test_resources.py.
+    # Generate the per-resource ResourceTestCase data driving test_resources.py.
     generated_test_cases.write_test_cases(output, schemas)
 
 

python/codegen/codegen/test_case.py.tmpl

@@ -1,12 +1,12 @@
 from databricks.bundles.core import Resources, ${singular}_mutator
 from databricks.bundles.core._generated.${plural} import _resource_type
-from databricks_tests.core._resource_test_case import TestCase
+from databricks_tests.core._resource_test_case import ResourceTestCase
 $model_imports
 
 
 def _test_case():
     return (
-        TestCase(
+        ResourceTestCase(
             add_resource=Resources.add_${singular},
             dict_example=$dict_example,
             dataclass_example=$dataclass_example,

python/codegen/codegen_tests/test_generated_dataclass.py

@@ -60,8 +60,6 @@ def test_generate_dataclass():
         extends=[],
         fields=[
             GeneratedField(
-                create_func_default=None,
-                create_func_type_name=variable_or_type(str_type(), is_required=True),
                 default=None,
                 default_factory=None,
                 description="task_key description",