sdk-python/scripts/gen_nexus_system_models.py at 629567c34d2974e16132c50840648b1f732e19ac · temporalio/sdk-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path

NEXUS_RPC_GEN_ENV_VAR = "TEMPORAL_NEXUS_RPC_GEN_DIR"
NEXUS_RPC_GEN_VERSION = "0.1.0-alpha.4"


def main() -> None:
    repo_root = Path(__file__).resolve().parent.parent
    # TODO: Remove the local .nexusrpc.yaml shim once the upstream API repo
    # checks in the Nexus definition we can consume directly.
    override_root = normalize_nexus_rpc_gen_root(
        Path.cwd(), env_value=NEXUS_RPC_GEN_ENV_VAR
    )
    input_schema = (
        repo_root
        / "temporalio"
        / "nexus"
        / "system"
        / "_workflow_service.nexusrpc.yaml"
    )
    output_file = (
        repo_root / "temporalio" / "nexus" / "system" / "_workflow_service_generated.py"
    )

    if not input_schema.is_file():
        raise RuntimeError(f"Expected Nexus schema at {input_schema}")

    run_nexus_rpc_gen(
        override_root=override_root,
        output_file=output_file,
        input_schema=input_schema,
    )
    subprocess.run(
        [
            "uv",
            "run",
            "ruff",
            "check",
            "--select",
            "I",
            "--fix",
            str(output_file),
        ],
        cwd=repo_root,
        check=True,
    )
    subprocess.run(
        [
            "uv",
            "run",
            "ruff",
            "format",
            str(output_file),
        ],
        cwd=repo_root,
        check=True,
    )


def run_nexus_rpc_gen(
    *, override_root: Path | None, output_file: Path, input_schema: Path
) -> None:
    common_args = [
        "--lang",
        "py",
        "--out-file",
        str(output_file),
        "--temporal-nexus-payload-codec-support",
        str(input_schema),
    ]
    if override_root is None:
        subprocess.run(
            ["npx", "--yes", f"nexus-rpc-gen@{NEXUS_RPC_GEN_VERSION}", *common_args],
            check=True,
        )
        return

    subprocess.run(
        [
            "node",
            "packages/nexus-rpc-gen/dist/index.js",
            *common_args,
        ],
        cwd=override_root,
        check=True,
    )


def normalize_nexus_rpc_gen_root(base_dir: Path, env_value: str) -> Path | None:
    raw_root = env_get(env_value)
    if raw_root is None:
        return None
    candidate = Path(raw_root)
    if not candidate.is_absolute():
        candidate = base_dir / candidate
    candidate = candidate.resolve()
    if (candidate / "package.json").is_file() and (candidate / "packages").is_dir():
        return candidate
    if (candidate / "src" / "package.json").is_file():
        return candidate / "src"
    raise RuntimeError(
        f"{NEXUS_RPC_GEN_ENV_VAR} must point to the nexus-rpc-gen repo root or its src directory"
    )


def env_get(name: str) -> str | None:
    return os.environ.get(name)


if __name__ == "__main__":
    try:
        main()
    except Exception as err:
        print(f"Failed to generate Nexus system models: {err}", file=sys.stderr)
        raise