Skip to content

Commit 483ee7d

Browse files
Sankalp-MittalIsaac
Sankalp-Mittal
and
Isaac
authored
Minor refactoring of PyDABs autogen code (#6575)
## Summary Minor follow-up refactoring on the PyDABs autogen code, addressing non-blocking review comments left on earlier PRs (#6421, #6528). ## Details - **Guard `generate_type` against unsupported map refs**: only `dict[str, str]` is modelled today; any other map ref (`map/int`, `map/<object>`) previously fell through and silently generated the bare value type instead of a dict. It now fails loudly, matching the existing guard in `generated_test_cases._synth_ref`. - **Remove the dead `create_func_type_name` / `create_func_default` fields** (and the now-orphaned `optional_type` helper): scaffolding for a static `create()` factory that was never implemented — set but never read since codegen was first added. - **De-duplicate `generate_field`**: collapse the four near-duplicate `GeneratedField(...)` branches into one base value plus per-branch `dataclasses.replace(...)` deltas. - **Rename `TestCase` → `ResourceTestCase`**: pytest no longer tries to collect it as a test class (its name no longer starts with `Test`), so the `__test__ = False` marker is unnecessary. The rename propagates through `test_case.py.tmpl` into `test_resources.py` and the generated per-resource files. - **Avoid a `KeyError` on an unknown launch stage**: use `_STAGE_RANK.get(prop.stage, 0)` instead of a bare dict lookup, so an absent/unknown stage is treated as GA rather than killing codegen (matches the Go `launchStageRank` semantics). - **De-duplicate module resolution**: route `_synth_ref`'s module lookup through `_module_of` to drop the duplicated `get_package(...) + assert` logic. All changes are behavior-preserving: regenerating the PyDABs output produces a byte-identical tree (aside from the `ResourceTestCase` rename propagating into the generated files). Codegen unit tests (38) and `test_resources.py` (232) pass. Supersedes #6573 (auto-closed when the stack it was based on merged to `main`). This pull request and its description were written by Isaac. --------- Co-authored-by: Isaac <no-reply@databricks.com>
1 parent e7d45a7 commit 483ee7d

32 files changed

Lines changed: 92 additions & 145 deletions

‎python/codegen/codegen/generated_dataclass.py‎

Lines changed: 24 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from dataclasses import dataclass
1+
from dataclasses import dataclass, replace
22
from typing import TYPE_CHECKING, Optional
33

44
if TYPE_CHECKING:
@@ -65,14 +65,6 @@ class GeneratedField:
6565
The type of the field in TypedDict, e.g., GeneratedType(name="TaskParam", ...)
6666
"""
6767

68-
create_func_type_name: GeneratedType
69-
"""
70-
Type type of the field in static "create" function, e.g., GeneratedType(name="TaskParam", ...)
71-
72-
It can be different from param_type_name because lists are made optional in "create" function
73-
to avoid problems with mutable default arguments.
74-
"""
75-
7668
description: Optional[str]
7769
"""
7870
The description of the field to be included into a docstring.
@@ -83,14 +75,6 @@ class GeneratedField:
8375
The default value of the field, e.g., "None"
8476
"""
8577

86-
create_func_default: Optional[str]
87-
"""
88-
The default value of the field in "create" function.
89-
90-
It can be different from default because lists are made optional in "create" function
91-
to avoid problems with mutable default arguments.
92-
"""
93-
9478
default_factory: Optional[str]
9579
"""
9680
Factory method for creating a default value, used for lists and dicts.
@@ -151,66 +135,27 @@ def generate_field(
151135
field_type = variable_or_type(field_type, is_required=is_required)
152136
param_type = variable_or_type(param_type, is_required=is_required)
153137

138+
# Base is the optional-scalar shape (default None). Collections instead use an
139+
# empty-container factory, and required fields carry no default.
140+
field = GeneratedField(
141+
field_name=field_name,
142+
type_name=field_type,
143+
param_type_name=param_type,
144+
description=prop.description,
145+
default="None",
146+
default_factory=None,
147+
experimental=is_experimental_stage(prop.stage),
148+
deprecated=prop.deprecated or False,
149+
)
150+
154151
if field_type.name == "VariableOrDict":
155-
return GeneratedField(
156-
field_name=field_name,
157-
type_name=field_type,
158-
param_type_name=param_type,
159-
create_func_type_name=optional_type(param_type),
160-
description=prop.description,
161-
default=None,
162-
default_factory="dict",
163-
create_func_default="None",
164-
experimental=is_experimental_stage(prop.stage),
165-
deprecated=prop.deprecated or False,
166-
)
152+
return replace(field, default=None, default_factory="dict")
167153
elif field_type.name == "VariableOrList":
168-
return GeneratedField(
169-
field_name=field_name,
170-
type_name=field_type,
171-
param_type_name=param_type,
172-
create_func_type_name=optional_type(param_type),
173-
description=prop.description,
174-
default=None,
175-
default_factory="list",
176-
create_func_default="None",
177-
experimental=is_experimental_stage(prop.stage),
178-
deprecated=prop.deprecated or False,
179-
)
154+
return replace(field, default=None, default_factory="list")
180155
elif is_required:
181-
return GeneratedField(
182-
field_name=field_name,
183-
type_name=field_type,
184-
param_type_name=param_type,
185-
create_func_type_name=param_type,
186-
description=prop.description,
187-
default=None,
188-
default_factory=None,
189-
create_func_default=None,
190-
experimental=is_experimental_stage(prop.stage),
191-
deprecated=prop.deprecated or False,
192-
)
193-
else:
194-
return GeneratedField(
195-
field_name=field_name,
196-
type_name=field_type,
197-
param_type_name=param_type,
198-
create_func_type_name=param_type,
199-
description=prop.description,
200-
default="None",
201-
default_factory=None,
202-
create_func_default="None",
203-
experimental=is_experimental_stage(prop.stage),
204-
deprecated=prop.deprecated or False,
205-
)
156+
return replace(field, default=None)
206157

207-
208-
def optional_type(generated: GeneratedType) -> GeneratedType:
209-
return GeneratedType(
210-
name="Optional",
211-
package="typing",
212-
parameters=[generated],
213-
)
158+
return field
214159

215160

216161
def str_type() -> GeneratedType:
@@ -281,7 +226,12 @@ def generate_type(namespace: str, ref: str, is_param: bool) -> GeneratedType:
281226
parameters=[element_type],
282227
)
283228

284-
if ref == "#/$defs/map/string":
229+
if ref.startswith("#/$defs/map/"):
230+
# Only dict[str, str] is modelled today; anything else (e.g. map/int,
231+
# map/<object>) fails loudly instead.
232+
if ref != "#/$defs/map/string":
233+
raise ValueError(f"Unsupported map ref: {ref}")
234+
285235
return dict_type()
286236

287237
class_name = packages.get_class_name(ref)

‎python/codegen/codegen/generated_test_cases.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
"""
2-
Generates the per-resource TestCase data driving databricks_tests/core/test_resources.py.
2+
Generates the per-resource ResourceTestCase data driving databricks_tests/core/test_resources.py.
33
44
For every wired resource a file _generated/<plural>.py is written (rendered from
5-
test_case.py.tmpl) exposing _test_case() -> (TestCase, _ResourceType). The generated
5+
test_case.py.tmpl) exposing _test_case() -> (ResourceTestCase, _ResourceType). The generated
66
_generated/__init__.py collects them into `test_cases`, which test_resources.py imports
77
and parametrizes its per-resource tests off.
88
@@ -160,8 +160,7 @@ def _synth_ref(
160160

161161
schema = schemas[name]
162162
class_name = packages.get_class_name(ref)
163-
module = packages.get_package(namespace, ref)
164-
assert module
163+
module = _module_of(namespace, ref)
165164

166165
if schema.type == openapi.SchemaType.STRING:
167166
value = schema.enum[0]
@@ -208,7 +207,7 @@ def _synth_object(
208207
continue
209208
if (
210209
prop.deprecated
211-
or _STAGE_RANK[prop.stage]
210+
or _STAGE_RANK.get(prop.stage, 0)
212211
> _STAGE_RANK[openapi.LaunchStage.PUBLIC_PREVIEW]
213212
):
214213
continue

‎python/codegen/codegen/main.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ def main(output: str):
5353
# decorators, and the core package __init__).
5454
generated_wiring.write_wiring(output)
5555

56-
# Generate the per-resource TestCase data driving test_resources.py.
56+
# Generate the per-resource ResourceTestCase data driving test_resources.py.
5757
generated_test_cases.write_test_cases(output, schemas)
5858

5959

‎python/codegen/codegen/test_case.py.tmpl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
from databricks.bundles.core import Resources, ${singular}_mutator
22
from databricks.bundles.core._generated.${plural} import _resource_type
3-
from databricks_tests.core._resource_test_case import TestCase
3+
from databricks_tests.core._resource_test_case import ResourceTestCase
44
$model_imports
55

66

77
def _test_case():
88
return (
9-
TestCase(
9+
ResourceTestCase(
1010
add_resource=Resources.add_${singular},
1111
dict_example=$dict_example,
1212
dataclass_example=$dataclass_example,

‎python/codegen/codegen_tests/test_generated_dataclass.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,6 @@ def test_generate_dataclass():
6060
extends=[],
6161
fields=[
6262
GeneratedField(
63-
create_func_default=None,
64-
create_func_type_name=variable_or_type(str_type(), is_required=True),
6563
default=None,
6664
default_factory=None,
6765
description="task_key description",

‎python/databricks_tests/core/_generated/alerts.py‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎python/databricks_tests/core/_generated/apps.py‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎python/databricks_tests/core/_generated/catalogs.py‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎python/databricks_tests/core/_generated/clusters.py‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎python/databricks_tests/core/_generated/database_catalogs.py‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)