Update system nexus generation to include notification service

changelog/added/chiming-jellyfish.md

@@ -0,0 +1 @@
+Add experimental notification service models in `temporalio.nexus.notifications`. The SDK does not currently use these models.

pyproject.toml

@@ -192,7 +192,7 @@ ignore_errors = true
 [tool.pydocstyle]
 convention = "google"
 # https://github.com/PyCQA/pydocstyle/issues/363#issuecomment-625563088
-match_dir = "^(?!(docs|scripts|tests|api|proto|system|\\.)).*"
+match_dir = "^(?!(docs|scripts|tests|api|proto|system|_support|notifications|\\.)).*"
 add_ignore = [
   # We like to wrap at a certain number of chars, even long summary sentences.
   # https://github.com/PyCQA/pydocstyle/issues/184

scripts/gen_nexus_system_api.py

@@ -1,11 +1,10 @@
+import ast
 import os
 import shutil
 import subprocess
 import sys
 import tempfile
-from importlib.util import module_from_spec, spec_from_file_location
 from pathlib import Path
-from typing import cast
 
 import gen_protos
 
@@ -22,26 +21,35 @@
     / "api_upstream"
     / "nexus"
 )
-wit_path = wit_input_dir / "workflow-service.wit"
 wit_deps_dir = wit_input_dir / "deps"
-output_dir = base_dir / "temporalio" / "nexus" / "system" / "workflow_service"
+workflow_output_dir = base_dir / "temporalio" / "nexus" / "system" / "workflow_service"
+notification_output_dir = base_dir / "temporalio" / "nexus" / "notifications"
 workflow_init_path = base_dir / "temporalio" / "workflow" / "__init__.py"
-workflowservice_request_response_proto = (
+NEX_GEN_VERSION = "0.2.8"
+proto_files = [
     gen_protos.api_proto_dir
     / "temporal"
     / "api"
     / "workflowservice"
     / "v1"
-    / "request_response.proto"
-)
-NEX_GEN_VERSION = "0.2.8"
+    / "request_response.proto",
+    gen_protos.api_proto_dir
+    / "temporal"
+    / "api"
+    / "notificationservice"
+    / "v1"
+    / "request_response.proto",
+]
 
 
 def nex_gen_command() -> list[str]:
     if bin_path := os.environ.get("NEX_GEN_BIN"):
         return [bin_path]
-
-    if shutil.which("nexgen") is None:
+    if (
+        shutil.which("nexgen") is None
+        or subprocess.check_output(["nexgen", "--version"], text=True).strip()
+        != f"nexgen {NEX_GEN_VERSION}"
+    ):
         subprocess.check_call(
             [
                 "cargo",
@@ -67,24 +75,65 @@ def build_descriptor_set(descriptor_path: Path) -> None:
             f"--proto_path={gen_protos.proto_dir}",
             "--include_imports",
             f"--descriptor_set_out={descriptor_path}",
-            str(workflowservice_request_response_proto),
+            *map(str, proto_files),
         ]
     )
 
 
-def generate_workflow_exports() -> None:
-    spec = spec_from_file_location(
-        "temporalio_nexus_system_workflow_service_exports",
-        output_dir / "__init__.py",
-        submodule_search_locations=[str(output_dir)],
+def generate_package(
+    command: list[str],
+    wit_name: str,
+    output_dir: Path,
+    descriptor_path: Path,
+    *,
+    native_api: bool,
+) -> None:
+    args = [*command, "python", str(wit_input_dir / wit_name), str(wit_deps_dir)]
+    if native_api:
+        args.append("--native-api")
+    subprocess.check_call(
+        [
+            *args,
+            "--system-nexus",
+            "--support-package",
+            "temporalio.nexus.system._support",
+            "--descriptors",
+            str(descriptor_path),
+            "--output",
+            str(output_dir),
+        ]
     )
-    if spec is None or spec.loader is None:
-        raise RuntimeError(f"Cannot load generated workflow service from {output_dir}")
-    module = module_from_spec(spec)
-    sys.modules[spec.name] = module
-    spec.loader.exec_module(module)
-    exports = cast(list[str], module.__all__)
 
+
+def generated_exports(init_path: Path) -> list[str]:
+    tree = ast.parse(init_path.read_text())
+    for statement in tree.body:
+        if not isinstance(statement, ast.Assign) or not any(
+            isinstance(target, ast.Name) and target.id == "__all__"
+            for target in statement.targets
+        ):
+            continue
+        value = ast.literal_eval(statement.value)
+        if not isinstance(value, list) or not all(
+            isinstance(item, str) for item in value
+        ):
+            raise RuntimeError(
+                f"generated package {init_path} __all__ must be a list of strings"
+            )
+        return value
+    raise RuntimeError(f"generated package {init_path} does not define __all__")
+
+
+def replace_marker_block(
+    content: str, begin: str, end: str, replacement: list[str]
+) -> str:
+    start = content.index(begin)
+    finish = content.index("\n", content.index(end, start)) + 1
+    return content[:start] + "".join(replacement) + content[finish:]
+
+
+def generate_workflow_exports() -> None:
+    exports = generated_exports(workflow_output_dir / "__init__.py")
     import_block = [
         "# BEGIN GENERATED NEXUS SYSTEM EXPORTS\n",
         "from temporalio.nexus.system.workflow_service import (\n",
@@ -98,71 +147,57 @@ def generate_workflow_exports() -> None:
         "    # END GENERATED NEXUS SYSTEM __ALL__\n",
     ]
     content = workflow_init_path.read_text()
-    start = content.index("# BEGIN GENERATED NEXUS SYSTEM EXPORTS")
-    end = content.index("# END GENERATED NEXUS SYSTEM EXPORTS", start)
-    end = content.index("\n", end) + 1
-    content = content[:start] + "".join(import_block) + content[end:]
-    start = content.index("    # BEGIN GENERATED NEXUS SYSTEM __ALL__")
-    end = content.index("    # END GENERATED NEXUS SYSTEM __ALL__", start)
-    end = content.index("\n", end) + 1
-    workflow_init_path.write_text(content[:start] + "".join(all_block) + content[end:])
+    content = replace_marker_block(
+        content, import_block[0].strip(), import_block[-1].strip(), import_block
+    )
+    workflow_init_path.write_text(
+        replace_marker_block(
+            content,
+            "    # BEGIN GENERATED NEXUS SYSTEM __ALL__",
+            "    # END GENERATED NEXUS SYSTEM __ALL__",
+            all_block,
+        )
+    )
 
 
 def generate_nexus_system_api() -> None:
-    if not wit_path.exists():
-        raise RuntimeError(f"missing WIT source: {wit_path}")
-    if not wit_deps_dir.exists():
-        raise RuntimeError(f"missing WIT dependency directory: {wit_deps_dir}")
-
+    required_paths = [
+        wit_input_dir / "workflow-service.wit",
+        wit_input_dir / "notification-service.wit",
+        wit_deps_dir,
+        *proto_files,
+    ]
+    for path in required_paths:
+        if not path.exists():
+            raise RuntimeError(f"missing generator input: {path}")
     with tempfile.TemporaryDirectory(dir=base_dir) as temp_dir:
         descriptor_path = Path(temp_dir) / "temporal_api.bin"
         build_descriptor_set(descriptor_path)
         command = nex_gen_command()
-
-        shutil.rmtree(output_dir, ignore_errors=True)
-        output_dir.parent.mkdir(parents=True, exist_ok=True)
-        subprocess.check_call(
-            [
-                *command,
-                "python",
-                str(wit_path),
-                str(wit_deps_dir),
-                "--native-api",
-                "--system-nexus",
-                "--support-package",
-                "temporalio.nexus.system._support",
-                "--descriptors",
-                str(descriptor_path),
-                "--output",
-                str(output_dir),
-            ]
-        )
-
-    (output_dir.parent / "__init__.py").touch()
+        for wit_name, output_dir, native_api in (
+            ("workflow-service.wit", workflow_output_dir, True),
+            ("notification-service.wit", notification_output_dir, False),
+        ):
+            shutil.rmtree(output_dir, ignore_errors=True)
+            output_dir.parent.mkdir(parents=True, exist_ok=True)
+            generate_package(
+                command,
+                wit_name,
+                output_dir,
+                descriptor_path,
+                native_api=native_api,
+            )
+    workflow_output_dir.parent.joinpath("__init__.py").touch()
     generate_workflow_exports()
+    format_paths = [
+        str(workflow_output_dir),
+        str(notification_output_dir),
+        str(workflow_init_path),
+    ]
     subprocess.check_call(
-        [
-            sys.executable,
-            "-m",
-            "ruff",
-            "check",
-            "--select",
-            "I",
-            "--fix",
-            str(output_dir),
-            str(workflow_init_path),
-        ]
-    )
-    subprocess.check_call(
-        [
-            sys.executable,
-            "-m",
-            "ruff",
-            "format",
-            str(output_dir),
-            str(workflow_init_path),
-        ]
+        [sys.executable, "-m", "ruff", "check", "--select", "I", "--fix", *format_paths]
     )
+    subprocess.check_call([sys.executable, "-m", "ruff", "format", *format_paths])
 
 
 if __name__ == "__main__":