PyDABs codegen: autogenerate core resource wiring

python/Taskfile.yml

@@ -76,6 +76,8 @@ tasks:
           ! -path databricks/bundles/core \
           ! -path databricks/bundles/resources \
           -exec rm -rf {} \;
+      # core/ is hand-written except for the generated wiring under _generated/.
+      - rm -rf databricks/bundles/core/_generated
       - cd codegen && uv run -m pytest codegen_tests
       - cd codegen && uv run -m codegen.main --output ..
       # Generated code is fixed and formatted by the global ruff (see ../ruff.toml).

python/codegen/codegen/.gitattributes

@@ -0,0 +1,2 @@
+# Render the wiring templates as Python on GitHub.
+*.py.tmpl linguist-language=Python

python/codegen/codegen/core_init.py.tmpl

@@ -0,0 +1,31 @@
+__all__ = [
+${all_block}
+]
+
+from databricks.bundles.core._bundle import Bundle
+from databricks.bundles.core._diagnostics import (
+    Diagnostic,
+    Diagnostics,
+    Severity,
+)
+from databricks.bundles.core._load import (
+    load_resources_from_current_package_module,
+    load_resources_from_module,
+    load_resources_from_modules,
+    load_resources_from_package_module,
+)
+from databricks.bundles.core._location import Location
+from databricks.bundles.core._resource import Resource
+from databricks.bundles.core._resource_mutator import ResourceMutator
+from databricks.bundles.core._resources import Resources
+from databricks.bundles.core._variable import (
+    Variable,
+    VariableOr,
+    VariableOrDict,
+    VariableOrList,
+    VariableOrOptional,
+    variables,
+)
+from databricks.bundles.core._generated import (
+${mutator_imports},
+)

python/codegen/codegen/generated_init.py.tmpl

@@ -0,0 +1,26 @@
+from typing import TYPE_CHECKING
+
+${imports}
+
+if TYPE_CHECKING:
+    from databricks.bundles.core._resource_type import _ResourceType
+
+__all__ = [
+${all_block}
+]
+
+
+class _GeneratedResources(
+${mixins}
+):
+    pass
+
+
+def _all_resource_types() -> "tuple[_ResourceType, ...]":
+    from databricks.bundles.core._generated import (
+${module_imports}
+    )
+
+    return (
+${type_entries}
+    )

python/codegen/codegen/generated_wiring.py

@@ -0,0 +1,160 @@
+"""
+Generates the per-resource wiring in databricks.bundles.core that used to be
+hand-written. All output is rendered from the *.py.tmpl templates in this
+directory.
+
+Each wired resource gets its own file _generated/<plural>.py (rendered from
+wiring_resource.py.tmpl): the add_* method + collection property mixin, the
+*_mutator decorator, and the _ResourceType entry. The generated
+_generated/__init__.py (generated_init.py.tmpl) collects them into
+_GeneratedResources (mixed into Resources), _all_resource_types(), and the
+mutator re-exports. The core package __init__ (core_init.py.tmpl) is generated
+too (static exports plus the generated mutators).
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+from string import Template
+
+import codegen.packages as packages
+
+HEADER = "# Code generated by pydabs-codegen. DO NOT EDIT.\n\n"
+
+
+def _load_template(name: str) -> Template:
+    return Template((Path(__file__).parent / name).read_text())
+
+
+_RESOURCE_TEMPLATE = _load_template("wiring_resource.py.tmpl")
+_GENERATED_INIT_TEMPLATE = _load_template("generated_init.py.tmpl")
+_CORE_INIT_TEMPLATE = _load_template("core_init.py.tmpl")
+
+
+@dataclass(frozen=True)
+class _WiredResource:
+    class_name: str
+    """Resource dataclass name, e.g. "Job"."""
+
+    singular_name: str
+    """Singular name used in methods and messages, e.g. "job"."""
+
+    plural_name: str
+    """Plural name, the same as the "resources" bundle section, e.g. "jobs"."""
+
+    model_module: str
+    """Module the resource dataclass lives in, e.g. "databricks.bundles.jobs._models.job"."""
+
+
+def _wired_resources() -> list[_WiredResource]:
+    # Every namespaced resource is wired. RESOURCE_NAMESPACE is the single source
+    # of truth for both model generation and wiring.
+    resources = []
+
+    for ref, namespace in packages.RESOURCE_NAMESPACE.items():
+        class_name = packages.get_class_name(ref)
+
+        resources.append(
+            _WiredResource(
+                class_name=class_name,
+                singular_name=class_name.lower(),
+                plural_name=namespace,
+                model_module=packages.get_package(namespace, ref),
+            )
+        )
+
+    resources.sort(key=lambda r: r.plural_name)
+
+    return resources
+
+
+def write_wiring(output: str):
+    resources = _wired_resources()
+
+    core_path = Path(output) / "databricks" / "bundles" / "core"
+    generated_path = core_path / "_generated"
+    generated_path.mkdir(parents=True, exist_ok=True)
+
+    for r in resources:
+        code = _RESOURCE_TEMPLATE.substitute(
+            {
+                "class": r.class_name,
+                "singular": r.singular_name,
+                "plural": r.plural_name,
+                "model_module": r.model_module,
+            }
+        )
+        (generated_path / f"{r.plural_name}.py").write_text(HEADER + code)
+
+    (generated_path / "__init__.py").write_text(HEADER + _collector_code(resources))
+    (core_path / "__init__.py").write_text(HEADER + _core_init_code(resources))
+
+    print(f"Writing wiring into {generated_path}")
+
+
+def _collector_code(resources: list[_WiredResource]) -> str:
+    imports = "\n".join(
+        f"from databricks.bundles.core._generated.{r.plural_name} import "
+        f"_{r.class_name}Resources, {r.singular_name}_mutator"
+        for r in resources
+    )
+
+    mutator_names = [f"{r.singular_name}_mutator" for r in resources]
+    exports = sorted(["_GeneratedResources", "_all_resource_types", *mutator_names])
+    all_block = "\n".join(f'    "{name}",' for name in exports)
+
+    mixins = "\n".join(f"    _{r.class_name}Resources," for r in resources)
+    module_imports = "\n".join(f"        {r.plural_name}," for r in resources)
+    type_entries = "\n".join(
+        f"        {r.plural_name}._resource_type()," for r in resources
+    )
+
+    return _GENERATED_INIT_TEMPLATE.substitute(
+        {
+            "imports": imports,
+            "all_block": all_block,
+            "mixins": mixins,
+            "module_imports": module_imports,
+            "type_entries": type_entries,
+        }
+    )
+
+
+# Static, non-resource exports of databricks.bundles.core, used to build the
+# sorted __all__ (the resource mutators are appended). Must stay in sync with the
+# import statements in core_init.py.tmpl.
+_CORE_INIT_STATIC_EXPORTS = [
+    "Bundle",
+    "Diagnostic",
+    "Diagnostics",
+    "Location",
+    "Resource",
+    "ResourceMutator",
+    "Resources",
+    "Severity",
+    "Variable",
+    "VariableOr",
+    "VariableOrDict",
+    "VariableOrList",
+    "VariableOrOptional",
+    "load_resources_from_current_package_module",
+    "load_resources_from_module",
+    "load_resources_from_modules",
+    "load_resources_from_package_module",
+    "variables",
+]
+
+
+def _core_init_code(resources: list[_WiredResource]) -> str:
+    mutator_names = [f"{r.singular_name}_mutator" for r in resources]
+
+    all_exports = sorted(_CORE_INIT_STATIC_EXPORTS + mutator_names)
+    all_block = "\n".join(f'    "{name}",' for name in all_exports)
+
+    mutator_imports = ",\n".join(f"    {name}" for name in sorted(mutator_names))
+
+    return _CORE_INIT_TEMPLATE.substitute(
+        {
+            "all_block": all_block,
+            "mutator_imports": mutator_imports,
+        }
+    )

python/codegen/codegen/main.py

@@ -8,6 +8,7 @@
 import codegen.generated_dataclass_patch as generated_dataclass_patch
 import codegen.generated_enum as generated_enum
 import codegen.generated_imports as generated_imports
+import codegen.generated_wiring as generated_wiring
 import codegen.jsonschema as openapi
 import codegen.jsonschema_patch as openapi_patch
 import codegen.packages as packages
@@ -46,6 +47,11 @@ def main(output: str):
 
         _write_exports(namespace, dataclasses, enums, output)
 
+    # Generate the per-resource wiring in databricks.bundles.core (the
+    # _ResourceType registry, Resources add_*/property methods, *_mutator
+    # decorators, and the core package __init__).
+    generated_wiring.write_wiring(output)
+
 
 def _transitively_mark_deprecated_and_private(
     roots: list[str],

python/codegen/codegen/wiring_resource.py.tmpl

@@ -0,0 +1,113 @@
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Optional, overload
+
+from databricks.bundles.core._bundle import Bundle
+from databricks.bundles.core._location import Location
+from databricks.bundles.core._resource_mutator import ResourceMutator
+from databricks.bundles.core._transform import _transform
+
+if TYPE_CHECKING:
+    from databricks.bundles.core._resource_type import _ResourceType
+    from ${model_module} import ${class}, ${class}Param
+
+
+def _resource_type() -> "_ResourceType":
+    from databricks.bundles.core._resource_type import _ResourceType
+    from ${model_module} import ${class}
+
+    return _ResourceType(
+        resource_type=${class},
+        singular_name="${singular}",
+        plural_name="${plural}",
+    )
+
+
+class _${class}Resources:
+    """
+    Generated ${singular} accessors, mixed into Resources.
+    """
+
+    # Provided by the Resources subclass; declared here so the generated methods
+    # below type-check.
+    _resources: dict[str, dict]
+
+    if TYPE_CHECKING:
+
+        def add_location(self, path: tuple[str, ...], location: Location) -> None: ...
+
+        def add_diagnostic_error(
+            self,
+            msg: str,
+            *,
+            detail: Optional[str] = None,
+            path: Optional[tuple[str, ...]] = None,
+            location: Optional[Location] = None,
+        ) -> None: ...
+
+    @property
+    def ${plural}(self) -> dict[str, "${class}"]:
+        return self._resources["${plural}"]
+
+    def add_${singular}(
+        self,
+        resource_name: str,
+        ${singular}: "${class}Param",
+        *,
+        location: Optional[Location] = None,
+    ) -> None:
+        """
+        Adds the resource ${singular} to the collection of resources. Resource name must be unique across all ${plural}.
+
+        :param resource_name: unique identifier for the ${singular}
+        :param ${singular}: the ${singular} to add, can be ${class} or dict
+        :param location: optional location of the ${singular} in the source code
+        """
+        from ${model_module} import ${class}
+
+        ${singular} = _transform(${class}, ${singular})
+        path = ("resources", "${plural}", resource_name)
+        location = location or Location.from_stack_frame(depth=1)
+
+        if self._resources["${plural}"].get(resource_name):
+            self.add_diagnostic_error(
+                msg=f"Duplicate resource name '{resource_name}' for resource '${singular}'. Resource names must be unique.",
+                location=location,
+                path=path,
+            )
+        else:
+            if location:
+                self.add_location(path, location)
+
+            self._resources["${plural}"][resource_name] = ${singular}
+
+
+@overload
+def ${singular}_mutator(
+    function: Callable[[Bundle, "${class}"], "${class}"],
+) -> ResourceMutator["${class}"]: ...
+
+
+@overload
+def ${singular}_mutator(
+    function: Callable[["${class}"], "${class}"],
+) -> ResourceMutator["${class}"]: ...
+
+
+def ${singular}_mutator(function: Callable) -> ResourceMutator["${class}"]:
+    """
+    Decorator for defining mutator for ${plural}. Function should return a new instance of the ${singular}
+    with the desired changes, instead of mutating the input ${singular}.
+
+    Example:
+
+    .. code-block:: python
+
+        @${singular}_mutator
+        def my_${singular}_mutator(bundle: Bundle, ${singular}: ${class}) -> ${class}:
+            return replace(${singular}, ...)
+
+    :param function: Function that mutates ${plural}.
+    """
+    from ${model_module} import ${class}
+
+    return ResourceMutator(resource_type=${class}, function=function)

python/databricks/bundles/.gitattributes

@@ -1,5 +1,6 @@
 # Generated by pydabs-codegen (see python/codegen). Each generated namespace
-# has a _models/ tree and an __init__.py; core/ is hand-written, so unset it.
+# has a _models/ tree and an __init__.py. In core/, only the _generated/ tree
+# and the package __init__.py are generated; the rest is hand-written.
 */_models/** linguist-generated=true
 */__init__.py linguist-generated=true
-core/__init__.py linguist-generated=false
+core/_generated/** linguist-generated=true