From ad29193ab0ac4572710b2bbd85165a0df1470b7f Mon Sep 17 00:00:00 2001 From: Sankalp-Mittal Date: Fri, 23 Oct 2026 11:07:56 +0000 Subject: [PATCH] [Python] Generalize recursive imports handling beyond jobs Detect dataclass cycles with a DFS from the resource root and quote the back-references, instead of hardcoding jobs.ForEachTask. Emit the _resolve_recursive_imports patch in any namespace that has quoted references, instead of only for jobs. Co-authored-by: Isaac --- .../codegen/generated_dataclass_patch.py | 48 +++++++++++++++--- python/codegen/codegen/main.py | 49 ++++++++++++------- .../test_generated_dataclass_patch.py | 44 +++++++++++++++++ python/databricks/bundles/jobs/__init__.py | 5 +- 4 files changed, 117 insertions(+), 29 deletions(-) create mode 100644 python/codegen/codegen_tests/test_generated_dataclass_patch.py diff --git a/python/codegen/codegen/generated_dataclass_patch.py b/python/codegen/codegen/generated_dataclass_patch.py index 9b2ae7ecba9..edf07a9c4e7 100644 --- a/python/codegen/codegen/generated_dataclass_patch.py +++ b/python/codegen/codegen/generated_dataclass_patch.py @@ -22,7 +22,7 @@ def reorder_required_fields(models: dict[str, GeneratedDataclass]): models[name] = replace(model, fields=required_fields + optional_fields) -def quote_recursive_references(models: dict[str, GeneratedDataclass]): +def quote_recursive_references(models: dict[str, GeneratedDataclass], root: str): """ If there is a cycle between two dataclasses, we need to quote one of them. @@ -33,15 +33,51 @@ class Foo: class Bar: foo: "Foo" + + Cycles are found with a DFS from the resource root, and the reference that + points back to a dataclass still being visited is quoted, e.g. ForEachTask.task + in jobs (Job -> Task -> ForEachTask -> Task). """ # see also _append_resolve_recursive_imports - if "jobs.ForEachTask" in models: - models["jobs.ForEachTask"] = _quote_recursive_references_for_model( - models["jobs.ForEachTask"], - references={"Task", "TaskParam"}, - ) + names_by_package = {model.package: name for name, model in models.items()} + visiting = set() + visited = set() + back_references: dict[str, set[str]] = {} + + def visit(name: str): + visiting.add(name) + + for field in models[name].fields: + for package in _referenced_packages(field.type_name): + # None for enums and core types, they can't form a cycle + ref = names_by_package.get(package) + + if ref in visiting: + model = models[ref] + back_references.setdefault(name, set()).update( + {model.class_name, f"{model.class_name}Param"} + ) + elif ref and ref not in visited: + visit(ref) + + visiting.remove(name) + visited.add(name) + + visit(root) + + for name, references in back_references.items(): + models[name] = _quote_recursive_references_for_model(models[name], references) + + +def _referenced_packages(type_name: GeneratedType) -> list[str]: + packages = [type_name.package] if type_name.package else [] + + for param in type_name.parameters: + packages += _referenced_packages(param) + + return packages def _quote_recursive_references_for_model( diff --git a/python/codegen/codegen/main.py b/python/codegen/codegen/main.py index 83735b8dc4c..13828d176e2 100644 --- a/python/codegen/codegen/main.py +++ b/python/codegen/codegen/main.py @@ -1,7 +1,6 @@ import argparse from dataclasses import replace from pathlib import Path -from textwrap import dedent import codegen.aliases_patch as aliases_patch import codegen.generated_dataclass as generated_dataclass @@ -44,7 +43,7 @@ def main(output: str): dataclasses, enums = _generate_code(namespace, reachable_schemas) generated_dataclass_patch.reorder_required_fields(dataclasses) - generated_dataclass_patch.quote_recursive_references(dataclasses) + generated_dataclass_patch.quote_recursive_references(dataclasses, root=resource) _write_code(dataclasses, enums, output) @@ -180,9 +179,7 @@ def _write_exports( generated_imports.append_dataclass_imports(b, dataclasses, exclude_packages=[]) generated_imports.append_enum_imports(b, enums, exclude_packages=[]) - # FIXME should be better generalized - if namespace == "jobs": - _append_resolve_recursive_imports(b) + _append_resolve_recursive_imports(b, dataclasses) if aliases: b.newline() @@ -199,28 +196,42 @@ def _write_exports( print(f"Writing exports into {source_path}") -def _append_resolve_recursive_imports(b: CodeBuilder): +def _append_resolve_recursive_imports( + b: CodeBuilder, + dataclasses: dict[str, GeneratedDataclass], +): """ Resolve forward references for recursive imports so we can assume that there are no forward references while inspecting type annotations. """ - b.append( - dedent(""" - def _resolve_recursive_imports(): - import typing + resolved = [] - from databricks.bundles.core._variable import VariableOr - from databricks.bundles.jobs._models.task import Task + for _, dataclass in sorted(dataclasses.items()): + imports = _collect_typechecking_imports(dataclass) + names = sorted({name for names in imports.values() for name in names}) - ForEachTask.__annotations__ = typing.get_type_hints( - ForEachTask, - globalns={"Task": Task, "VariableOr": VariableOr}, - ) + if not names: + continue - _resolve_recursive_imports() - """) - ) + # quoted classes are imported at runtime by this __init__.py, see quote_recursive_references + globalns = ", ".join(f'"{name}": {name}' for name in names) + resolved.append( + f" {dataclass.class_name}.__annotations__ = typing.get_type_hints(\n" + f" {dataclass.class_name},\n" + f" globalns={{{globalns}}},\n" + f" )\n" + ) + + if not resolved: + return + + b.newline() + b.append("def _resolve_recursive_imports():\n") + b.append(" import typing\n\n") + b.append("\n".join(resolved)) + b.newline() + b.append("_resolve_recursive_imports()").newline() def _collect_typechecking_imports( diff --git a/python/codegen/codegen_tests/test_generated_dataclass_patch.py b/python/codegen/codegen_tests/test_generated_dataclass_patch.py new file mode 100644 index 00000000000..f738721c0e6 --- /dev/null +++ b/python/codegen/codegen_tests/test_generated_dataclass_patch.py @@ -0,0 +1,44 @@ +from codegen.generated_dataclass import generate_dataclass +from codegen.generated_dataclass_patch import quote_recursive_references +from codegen.jsonschema import Property, Schema, SchemaType + + +def _model(schema_name: str, properties: dict[str, str]): + return generate_dataclass( + namespace="bananas", + schema_name=schema_name, + schema=Schema( + type=SchemaType.OBJECT, + properties={name: Property(ref=ref) for name, ref in properties.items()}, + ), + ) + + +def _type_names(model) -> set[str]: + def visit(type_name): + return {type_name.name}.union(*(visit(p) for p in type_name.parameters)) + + out = set() + for field in model.fields: + out |= visit(field.type_name) | visit(field.param_type_name) + return out + + +def test_quote_recursive_references(): + # Tree -> Node -> Branch -> list[Node]: only the reference back to Node, + # which is still being visited from the root, gets quoted + models = { + "bananas.Tree": _model("bananas.Tree", {"node": "#/$defs/bananas.Node"}), + "bananas.Node": _model("bananas.Node", {"branch": "#/$defs/bananas.Branch"}), + "bananas.Branch": _model( + "bananas.Branch", + {"nodes": "#/$defs/slice/bananas.Node", "tree_name": "#/$defs/string"}, + ), + } + + quote_recursive_references(models, root="bananas.Tree") + + assert '"Node"' in _type_names(models["bananas.Branch"]) + assert '"NodeParam"' in _type_names(models["bananas.Branch"]) + assert "Node" in _type_names(models["bananas.Tree"]) + assert "Branch" in _type_names(models["bananas.Node"]) diff --git a/python/databricks/bundles/jobs/__init__.py b/python/databricks/bundles/jobs/__init__.py index 1df8be62977..a729caa11c6 100644 --- a/python/databricks/bundles/jobs/__init__.py +++ b/python/databricks/bundles/jobs/__init__.py @@ -927,12 +927,9 @@ def _resolve_recursive_imports(): import typing - from databricks.bundles.core._variable import VariableOr - from databricks.bundles.jobs._models.task import Task - ForEachTask.__annotations__ = typing.get_type_hints( ForEachTask, - globalns={"Task": Task, "VariableOr": VariableOr}, + globalns={"Task": Task, "TaskParam": TaskParam}, )