Early draft integration of nex-gen into python

Author: tconley1428

pyproject.toml

@@ -98,14 +98,17 @@ format = [
   { cmd = "cargo fmt", cwd = "temporalio/bridge" },
 ]
 gen-docs = "uv run scripts/gen_docs.py"
+gen-nexus-system-api = "uv run scripts/gen_nexus_system_api.py"
 gen-protos = [
   { cmd = "uv run scripts/gen_protos.py" },
+  { ref = "gen-nexus-system-api" },
   { cmd = "uv run scripts/gen_payload_visitor.py" },
   { cmd = "uv run scripts/gen_bridge_client.py" },
   { ref = "format" },
 ]
 gen-protos-docker = [
   { cmd = "uv run scripts/gen_protos_docker.py" },
+  { ref = "gen-nexus-system-api" },
   { cmd = "uv run scripts/gen_payload_visitor.py" },
   { cmd = "uv run scripts/gen_bridge_client.py" },
   { ref = "format" },
@@ -169,7 +172,7 @@ exclude = [
 [tool.pydocstyle]
 convention = "google"
 # https://github.com/PyCQA/pydocstyle/issues/363#issuecomment-625563088
-match_dir = "^(?!(docs|scripts|tests|api|proto|\\.)).*"
+match_dir = "^(?!(docs|scripts|tests|api|proto|_system|\\.)).*"
 add_ignore = [
   # We like to wrap at a certain number of chars, even long summary sentences.
   # https://github.com/PyCQA/pydocstyle/issues/184

scripts/_nexus/temporal-system.wit

@@ -0,0 +1,76 @@
+package temporal:nexus@1.0.0;
+
+world system {
+  export workflow-service;
+}
+
+/// @nexus.endpoint "temporal-system"
+interface workflow-service {
+  use nexus:temporal-types/model@1.0.0.{
+    duration,
+    memo,
+    payloads,
+    placeholder,
+    priority,
+    retry-policy,
+    search-attributes,
+    signal-function,
+    task-queue,
+    user-metadata,
+    versioning-override,
+    workflow-function,
+    workflow-id-conflict-policy,
+    workflow-id-reuse-policy,
+  };
+
+  /// @nexus.proto "temporal.api.workflowservice.v1.SignalWithStartWorkflowExecutionRequest"
+  record signal-with-start-workflow-execution-request {
+    /// @nexus.proto-field "workflow_type"
+    workflow: workflow-function,
+    workflow-id: string,
+    task-queue: task-queue,
+    /// @nexus.proto-field "signal_name"
+    signal: signal-function,
+    workflow-execution-timeout: option<duration>,
+    workflow-run-timeout: option<duration>,
+    workflow-task-timeout: option<duration>,
+    identity: option<string>,
+    request-id: option<string>,
+    workflow-id-reuse-policy: option<workflow-id-reuse-policy>,
+    workflow-id-conflict-policy: option<workflow-id-conflict-policy>,
+    retry-policy: option<retry-policy>,
+    cron-schedule: option<string>,
+    memo: option<memo>,
+    search-attributes: option<search-attributes>,
+    priority: option<priority>,
+    versioning-override: option<versioning-override>,
+    workflow-start-delay: option<duration>,
+    user-metadata: option<user-metadata>,
+    /// @nexus.source python="workflow_namespace" typescript="workflowNamespace"
+    namespace: string,
+    /// @nexus.omit
+    control: placeholder,
+    /// @nexus.omit
+    header: placeholder,
+    /// @nexus.omit
+    links: placeholder,
+    /// @nexus.omit
+    time-skipping-config: placeholder,
+  }
+
+  /// @nexus.proto "temporal.api.workflowservice.v1.SignalWithStartWorkflowExecutionResponse"
+  record signal-with-start-workflow-execution-response {
+    run-id: option<string>,
+    started: option<bool>,
+    /// @nexus.omit
+    signal-link: placeholder,
+  }
+
+  /// @nexus.output-transform
+  ///   python-type="workflow.ExternalWorkflowHandle[typing.Any]"
+  ///   python="workflow.get_external_workflow_handle(request.workflow_id, run_id=result.run_id)"
+  /// @nexus.operation name="SignalWithStartWorkflowExecution"
+  signal-with-start-workflow: func(
+    request: signal-with-start-workflow-execution-request,
+  ) -> signal-with-start-workflow-execution-response;
+}

scripts/gen_nexus_system_api.py

@@ -0,0 +1,123 @@
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import gen_protos
+
+base_dir = Path(__file__).parent.parent
+wit_path = base_dir / "scripts" / "_nexus" / "temporal-system.wit"
+output_dir = base_dir / "temporalio" / "nexus" / "_system" / "workflow_service"
+workflowservice_request_response_proto = (
+    gen_protos.api_proto_dir
+    / "temporal"
+    / "api"
+    / "workflowservice"
+    / "v1"
+    / "request_response.proto"
+)
+
+
+def nexus_api_gen_command() -> list[str]:
+    if bin_path := os.environ.get("NEXUS_API_GEN_BIN"):
+        return [bin_path]
+
+    candidates = []
+    if source_dir := os.environ.get("NEXUS_API_GEN_DIR"):
+        candidates.append(Path(source_dir))
+    candidates.extend(
+        [
+            base_dir / "tools" / "nexus-api-gen",
+            base_dir.parent / "nexus-api-gen",
+        ]
+    )
+    for candidate in candidates:
+        manifest_path = candidate / "Cargo.toml"
+        if manifest_path.exists():
+            return ["cargo", "run", "--manifest-path", str(manifest_path), "--"]
+
+    raise RuntimeError(
+        "nexus-api-gen not found; set NEXUS_API_GEN_BIN or NEXUS_API_GEN_DIR"
+    )
+
+
+def build_descriptor_set(descriptor_path: Path) -> None:
+    subprocess.check_call(
+        [
+            sys.executable,
+            "-mgrpc_tools.protoc",
+            f"--proto_path={gen_protos.api_proto_dir}",
+            f"--proto_path={gen_protos.proto_dir}",
+            "--include_imports",
+            f"--descriptor_set_out={descriptor_path}",
+            str(workflowservice_request_response_proto),
+        ]
+    )
+
+
+def strip_unsupported_pyright_comments() -> None:
+    for path in output_dir.rglob("*.py"):
+        content = path.read_text()
+        content = content.replace("# pyright: reportAny=false\n", "")
+        content = content.replace(
+            "# pyright: reportAny=false, reportExplicitAny=false\n", ""
+        )
+        path.write_text(content)
+
+
+def generate_nexus_system_api() -> None:
+    if not wit_path.exists():
+        raise RuntimeError(f"missing WIT source: {wit_path}")
+
+    with tempfile.TemporaryDirectory(dir=base_dir) as temp_dir:
+        descriptor_path = Path(temp_dir) / "temporal_api.bin"
+        build_descriptor_set(descriptor_path)
+
+        shutil.rmtree(output_dir, ignore_errors=True)
+        output_dir.parent.mkdir(parents=True, exist_ok=True)
+        subprocess.check_call(
+            [
+                *nexus_api_gen_command(),
+                "generate",
+                "--lang",
+                "python",
+                "--input",
+                str(wit_path),
+                "--descriptors",
+                str(descriptor_path),
+                "--output",
+                str(output_dir),
+            ]
+        )
+
+    (output_dir.parent / "__init__.py").touch()
+    strip_unsupported_pyright_comments()
+    subprocess.check_call(
+        [
+            sys.executable,
+            "-m",
+            "ruff",
+            "check",
+            "--select",
+            "I",
+            "--fix",
+            str(output_dir),
+        ]
+    )
+    subprocess.check_call(
+        [
+            sys.executable,
+            "-m",
+            "ruff",
+            "format",
+            str(output_dir),
+        ]
+    )
+
+
+if __name__ == "__main__":
+    print("Generating Nexus system API...", file=sys.stderr)
+    generate_nexus_system_api()
+    print("Done", file=sys.stderr)

temporalio/nexus/_system/__init__.py

@@ -0,0 +1,21 @@
+# Generated by nexus-api-gen. DO NOT EDIT!
+
+from __future__ import annotations
+
+from . import models
+from . import service as _service
+from .operations.signal_with_start_workflow import signal_with_start_workflow
+
+__all__ = [
+    "models",
+    "signal_with_start_workflow",
+    "__nexus_operation_registry__",
+]
+
+
+__nexus_operation_registry__ = {
+    (
+        "WorkflowService",
+        "SignalWithStartWorkflowExecution",
+    ): _service.WorkflowService.signal_with_start_workflow,
+}

temporalio/nexus/_system/workflow_service/__init__.py

@@ -0,0 +1,5 @@
+# Generated by nexus-api-gen. DO NOT EDIT!
+
+from __future__ import annotations
+
+__all__ = []

temporalio/nexus/_system/workflow_service/_resources/__init__.py

@@ -0,0 +1,5 @@
+# Generated by nexus-api-gen. DO NOT EDIT!
+
+from __future__ import annotations
+
+from .model_overrides import *  # noqa: F401,F403