|
| 1 | +import os |
| 2 | +import shutil |
| 3 | +import subprocess |
| 4 | +import sys |
| 5 | +import tempfile |
| 6 | +from importlib.util import module_from_spec, spec_from_file_location |
| 7 | +from pathlib import Path |
| 8 | +from typing import cast |
| 9 | + |
| 10 | +import gen_protos |
| 11 | + |
| 12 | +base_dir = Path(__file__).parent.parent |
| 13 | +sys.path.insert(0, str(base_dir)) |
| 14 | +wit_input_dir = ( |
| 15 | + base_dir |
| 16 | + / "temporalio" |
| 17 | + / "bridge" |
| 18 | + / "sdk-core" |
| 19 | + / "crates" |
| 20 | + / "protos" |
| 21 | + / "protos" |
| 22 | + / "api_upstream" |
| 23 | + / "nexus" |
| 24 | +) |
| 25 | +wit_path = wit_input_dir / "workflow-service.wit" |
| 26 | +wit_deps_dir = wit_input_dir / "deps" |
| 27 | +python_support_path = base_dir / "scripts" / "nex_gen_support.py" |
| 28 | +output_dir = base_dir / "temporalio" / "nexus" / "system" / "workflow_service" |
| 29 | +workflow_init_path = base_dir / "temporalio" / "workflow" / "__init__.py" |
| 30 | +workflowservice_request_response_proto = ( |
| 31 | + gen_protos.api_proto_dir |
| 32 | + / "temporal" |
| 33 | + / "api" |
| 34 | + / "workflowservice" |
| 35 | + / "v1" |
| 36 | + / "request_response.proto" |
| 37 | +) |
| 38 | + |
| 39 | + |
| 40 | +def nex_gen_command() -> list[str]: |
| 41 | + if bin_path := os.environ.get("NEX_GEN_BIN"): |
| 42 | + return [bin_path] |
| 43 | + |
| 44 | + if shutil.which("nex-gen") is None: |
| 45 | + subprocess.check_call(["cargo", "install", "--locked", "nex-gen", "--force"]) |
| 46 | + return ["nex-gen"] |
| 47 | + |
| 48 | + |
| 49 | +def build_descriptor_set(descriptor_path: Path) -> None: |
| 50 | + subprocess.check_call( |
| 51 | + [ |
| 52 | + sys.executable, |
| 53 | + "-mgrpc_tools.protoc", |
| 54 | + f"--proto_path={gen_protos.api_proto_dir}", |
| 55 | + f"--proto_path={gen_protos.proto_dir}", |
| 56 | + "--include_imports", |
| 57 | + f"--descriptor_set_out={descriptor_path}", |
| 58 | + str(workflowservice_request_response_proto), |
| 59 | + ] |
| 60 | + ) |
| 61 | + |
| 62 | + |
| 63 | +def generate_workflow_exports() -> None: |
| 64 | + spec = spec_from_file_location( |
| 65 | + "temporalio_nexus_system_workflow_service_exports", |
| 66 | + output_dir / "__init__.py", |
| 67 | + submodule_search_locations=[str(output_dir)], |
| 68 | + ) |
| 69 | + if spec is None or spec.loader is None: |
| 70 | + raise RuntimeError(f"Cannot load generated workflow service from {output_dir}") |
| 71 | + module = module_from_spec(spec) |
| 72 | + sys.modules[spec.name] = module |
| 73 | + spec.loader.exec_module(module) |
| 74 | + exports = cast(list[str], module.__all__) |
| 75 | + |
| 76 | + import_block = [ |
| 77 | + "# BEGIN GENERATED NEXUS SYSTEM EXPORTS\n", |
| 78 | + "from temporalio.nexus.system.workflow_service import (\n", |
| 79 | + *[f" {export},\n" for export in exports], |
| 80 | + ")\n", |
| 81 | + "# END GENERATED NEXUS SYSTEM EXPORTS\n", |
| 82 | + ] |
| 83 | + all_block = [ |
| 84 | + " # BEGIN GENERATED NEXUS SYSTEM __ALL__\n", |
| 85 | + *[f' "{export}",\n' for export in exports], |
| 86 | + " # END GENERATED NEXUS SYSTEM __ALL__\n", |
| 87 | + ] |
| 88 | + content = workflow_init_path.read_text() |
| 89 | + start = content.index("# BEGIN GENERATED NEXUS SYSTEM EXPORTS") |
| 90 | + end = content.index("# END GENERATED NEXUS SYSTEM EXPORTS", start) |
| 91 | + end = content.index("\n", end) + 1 |
| 92 | + content = content[:start] + "".join(import_block) + content[end:] |
| 93 | + start = content.index(" # BEGIN GENERATED NEXUS SYSTEM __ALL__") |
| 94 | + end = content.index(" # END GENERATED NEXUS SYSTEM __ALL__", start) |
| 95 | + end = content.index("\n", end) + 1 |
| 96 | + workflow_init_path.write_text(content[:start] + "".join(all_block) + content[end:]) |
| 97 | + |
| 98 | + |
| 99 | +def generate_nexus_system_api() -> None: |
| 100 | + if not wit_path.exists(): |
| 101 | + raise RuntimeError(f"missing WIT source: {wit_path}") |
| 102 | + if not wit_deps_dir.exists(): |
| 103 | + raise RuntimeError(f"missing WIT dependency directory: {wit_deps_dir}") |
| 104 | + if not python_support_path.exists(): |
| 105 | + raise RuntimeError(f"missing Python support source: {python_support_path}") |
| 106 | + |
| 107 | + with tempfile.TemporaryDirectory(dir=base_dir) as temp_dir: |
| 108 | + descriptor_path = Path(temp_dir) / "temporal_api.bin" |
| 109 | + build_descriptor_set(descriptor_path) |
| 110 | + command = nex_gen_command() |
| 111 | + |
| 112 | + shutil.rmtree(output_dir, ignore_errors=True) |
| 113 | + output_dir.parent.mkdir(parents=True, exist_ok=True) |
| 114 | + subprocess.check_call( |
| 115 | + [ |
| 116 | + *command, |
| 117 | + "generate", |
| 118 | + "--lang", |
| 119 | + "python", |
| 120 | + "--input", |
| 121 | + str(wit_path), |
| 122 | + "--input", |
| 123 | + str(wit_deps_dir), |
| 124 | + "--support-file", |
| 125 | + str(python_support_path), |
| 126 | + "--descriptors", |
| 127 | + str(descriptor_path), |
| 128 | + "--output", |
| 129 | + str(output_dir), |
| 130 | + ] |
| 131 | + ) |
| 132 | + |
| 133 | + (output_dir.parent / "__init__.py").touch() |
| 134 | + generate_workflow_exports() |
| 135 | + subprocess.check_call( |
| 136 | + [ |
| 137 | + sys.executable, |
| 138 | + "-m", |
| 139 | + "ruff", |
| 140 | + "check", |
| 141 | + "--select", |
| 142 | + "I", |
| 143 | + "--fix", |
| 144 | + str(output_dir), |
| 145 | + str(workflow_init_path), |
| 146 | + ] |
| 147 | + ) |
| 148 | + subprocess.check_call( |
| 149 | + [ |
| 150 | + sys.executable, |
| 151 | + "-m", |
| 152 | + "ruff", |
| 153 | + "format", |
| 154 | + str(output_dir), |
| 155 | + str(workflow_init_path), |
| 156 | + ] |
| 157 | + ) |
| 158 | + |
| 159 | + |
| 160 | +if __name__ == "__main__": |
| 161 | + print("Generating Nexus system API...", file=sys.stderr) |
| 162 | + generate_nexus_system_api() |
| 163 | + print("Done", file=sys.stderr) |
0 commit comments