[Python] Generalize recursive imports handling beyond jobs

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(

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(

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"])