1
import argparse
2
import json
3
import re
4
import shutil
5
import subprocess
6
import sys
7
from copy import deepcopy
8
from functools import cache
9
from os import environ
10
from pathlib import Path
11
from tempfile import TemporaryDirectory
12
from textwrap import dedent, indent
13
 
14
from tool_utils import hexdigest, run, timing
15
 
16
 
17
def cquery(rule):
18
    res = subprocess.run(
19
        [
20
            "bazel",
21
            "cquery",
22
            rule,
23
            "--output=files",
24
        ],
25
        capture_output=True,
26
        text=True,
27
        check=False,
28
    )
29
    if res.returncode:
30
        print(res.stdout)
31
        print(res.stderr)
32
        sys.exit(res.returncode)
33
    return res.stdout.strip()
34
 
35
 
36
@cache
37
def _bundle_version_info():
38
    with Path(cquery("@workerd//src/pyodide:bundle_version_info")).open() as f:
39
        return json.load(f)
40
 
41
 
42
def bundle_version_info():
43
    return deepcopy(_bundle_version_info())
44
 
45
 
46
TEMPLATE = """
47
using Workerd = import "/workerd/workerd.capnp";
48
49
const config :Workerd.Config = (
50
  services = [
51
    (name = "main", worker = .mainWorker),
52
  ],
53
);
54
55
const mainWorker :Workerd.Worker = (
56
  modules = [
57
    (name = "worker.py", pythonModule = embed "./worker.py"),
58
    {requirements}
59
  ],
60
  compatibilityDate = "2025-08-05",
61
  compatibilityFlags = ["python_no_global_handlers", {compat_flags}],
62
  # Learn more about compatibility dates at:
63
  # https://developers.cloudflare.com/workers/platform/compatibility-dates/
64
);
65
"""
66
 
67
 
68
def make_config(
69
    flags: list[str],
70
    reqs: list[str],
71
) -> str:
72
    requirements = ""
73
    for name in reqs:
74
        requirements += f'(name="{name}", pythonRequirement=""),'
75
 
76
    compat_flags = ""
77
    for flag in flags:
78
        compat_flags += f'"{flag}", '
79
    return TEMPLATE.format(requirements=requirements, compat_flags=compat_flags)
80
 
81
 
82
def make_worker(imports: list[str]) -> str:
83
    contents = ""
84
    for i in imports:
85
        contents += f"import {i}\n"
86
    contents += dedent("""\
87
    from workers import WorkerEntrypoint
88
    class Default(WorkerEntrypoint):
89
        def test(self):
90
            pass
91
    """)
92
    return contents
93
 
94
 
95
def make_snapshot(  # noqa: PLR0913
96
    d: Path,
97
    outdir: Path,
98
    outprefix: str,
99
    compat_flags: list[str],
100
    requirements: list[str],
101
    imports: list[str],
102
) -> str:
103
    config_path = d / "config.capnp"
104
    config_path.write_text(make_config(compat_flags, requirements))
105
    worker_path = d / "worker.py"
106
    worker_path.write_text(make_worker(imports))
107
    if imports:
108
        snapshot_flag = "--python-save-snapshot"
109
    else:
110
        snapshot_flag = "--python-save-baseline-snapshot"
111
 
112
    if "WORKERD_BINARY" in environ:
113
        workerd = [environ["WORKERD_BINARY"]]
114
    else:
115
        workerd = [
116
            "bazel",
117
            "run",
118
            "@workerd//src/workerd/server:workerd",
119
            "--",
120
        ]
121
    run(
122
        [
123
            *workerd,
124
            "test",
125
            config_path,
126
            snapshot_flag,
127
            "--pyodide-bundle-disk-cache-dir",
128
            d,
129
            "--pyodide-package-disk-cache-dir",
130
            d,
131
            "--experimental",
132
        ],
133
    )
134
    snapshot_path = d / "snapshot.bin"
135
    digest = hexdigest(snapshot_path)
136
    digest9 = digest[:9]
137
    outname = f"{outprefix}-{digest9}.bin"
138
    outfile = outdir / outname
139
    shutil.copyfile(snapshot_path, outfile)
140
    snapshot_path.unlink()
141
    return [outname, digest]
142
 
143
 
144
def make_baseline_snapshot(
145
    cache: Path, outdir: Path, compat_flags: list[str]
146
) -> list[tuple[str, str]]:
147
    name, digest = make_snapshot(cache, outdir, "baseline", compat_flags, [], [])
148
    return [
149
        ("baseline_snapshot", name),
150
        ("baseline_snapshot_hash", digest),
151
    ]
152
 
153
 
154
def make_numpy_snapshot(
155
    cache: Path, outdir: Path, compat_flags: list[str]
156
) -> list[tuple[str, str]]:
157
    name, digest = make_snapshot(
158
        cache, outdir, "package_snapshot_numpy", compat_flags, ["numpy"], ["numpy"]
159
    )
160
    return [
161
        ("numpy_snapshot", name),
162
        ("numpy_snapshot_hash", digest),
163
    ]
164
 
165
 
166
def make_fastapi_snapshot(
167
    cache: Path, outdir: Path, compat_flags: list[str]
168
) -> list[tuple[str, str]]:
169
    name, digest = make_snapshot(
170
        cache,
171
        outdir,
172
        "package_snapshot_fastapi",
173
        compat_flags,
174
        ["fastapi"],
175
        ["fastapi", "pydantic"],
176
    )
177
    return [
178
        ("fastapi_snapshot", name),
179
        ("fastapi_snapshot_hash", digest),
180
    ]
181
 
182
 
183
def make_snapshots(
184
    cache: Path, outdir: Path, update_released: bool
185
) -> tuple[str, tuple[str, str]]:
186
    res = []
187
    for ver, info in bundle_version_info().items():
188
        if ver.startswith("dev"):
189
            continue
190
        if not update_released and info.get("released", False):
191
            continue
192
        compat_flags = list({"python_workers", info["enable_flag_name"]})
193
 
194
        ver_info = []
195
        with timing(f"version {ver} snapshots"):
196
            with timing("baseline snapshot"):
197
                ver_info += make_baseline_snapshot(cache, outdir, compat_flags)
198
            with timing("numpy snapshot"):
199
                ver_info += make_numpy_snapshot(cache, outdir, compat_flags)
200
            with timing("fastapi snapshot"):
201
                ver_info += make_fastapi_snapshot(cache, outdir, compat_flags)
202
        res.append((ver, ver_info))
203
    return res
204
 
205
 
206
def update_python_metadata_bzl(res: tuple[str, tuple[str, str]]):
207
    """Update python_metadata.bzl file with new snapshot values."""
208
    metadata_path = (
209
        Path(__file__).parent.parent.parent / "build" / "python_metadata.bzl"
210
    )
211
    content = metadata_path.read_text()
212
 
213
    for ver, kvs in res:
214
        # Find the version block and update snapshot values
215
        version_pattern = rf'(\s+{{\s*\n\s*"name":\s*"{re.escape(ver)}",.*?)}}'
216
 
217
        def replace_version_block(match, *, kvs=kvs):
218
            block = match.group(1)
219
            # Update each key-value pair
220
            for key, val in kvs:
221
                key_pattern = rf'("{re.escape(key)}":\s*)"[^"]*"'
222
                block = re.sub(key_pattern, rf'\1"{val}"', block)
223
            return block + "}"
224
 
225
        content = re.sub(
226
            version_pattern, replace_version_block, content, flags=re.DOTALL
227
        )
228
 
229
    metadata_path.write_text(content)
230
 
231
 
232
def upload_snapshots(outdir: Path):
233
    from boto3 import client
234
 
235
    s3 = client(
236
        "s3",
237
        endpoint_url=f"https://{environ['R2_ACCOUNT_ID']}.r2.cloudflarestorage.com",
238
        aws_access_key_id=environ["R2_ACCESS_KEY_ID"],
239
        aws_secret_access_key=environ["R2_SECRET_ACCESS_KEY"],
240
        region_name="auto",
241
    )
242
 
243
    for file in outdir.glob("*.bin"):
244
        if file.name.startswith("baseline-"):
245
            key = "baseline-snapshot/" + hexdigest(file)
246
        else:
247
            key = "test-snapshot/" + file.name
248
        s3.upload_file(str(file), "pyodide-capnp-bin", key)
249
 
250
 
251
def main() -> int:
252
    parser = argparse.ArgumentParser(
253
        description="Upload Pyodide bundles and update metadata"
254
    )
255
    parser.add_argument(
256
        "--update-released",
257
        action="store_true",
258
        help="Update already released versions?",
259
    )
260
    args = parser.parse_args()
261
 
262
    subprocess.run(
263
        [
264
            "bazel",
265
            "build",
266
            "@workerd//src/pyodide:bundle_version_info",
267
        ],
268
        check=True,
269
    )
270
 
271
    # Create generated-snapshots directory
272
    outdir = Path(__file__).parent / "generated-snapshots"
273
    if outdir.exists() and outdir.is_dir() and any(outdir.iterdir()):
274
        print(f"Error: Directory {outdir} exists and is not empty", file=sys.stderr)
275
        return 1
276
    outdir.mkdir(parents=True, exist_ok=True)
277
 
278
    with TemporaryDirectory() as package_cache:
279
        cache = Path(package_cache)
280
        res = make_snapshots(cache, outdir, args.update_released)
281
 
282
    update_python_metadata_bzl(res)
283
 
284
    upload_snapshots(outdir)
285
    print()
286
    print(
287
        "Upload these files to the ew-snapshot-tests R2 bucket: "
288
        + "https://dash.cloudflare.com/e415f1017791ced9d5f3eb0df2b31c9e/r2/default/buckets/ew-snapshot-tests"
289
    )
290
    print("Updated python_metadata.bzl with:")
291
    for ver, kvs in res:
292
        print("Version", ver)
293
        for key, val in kvs:
294
            print(indent(f'"{key}": "{val}",', " " * 8))
295
    return 0
296
 
297
 
298
if __name__ == "__main__":
299
    sys.exit(main())