1
#!/usr/bin/env python3
2
"""
3
Script to vendor a Python package for vendored_py_wd_test.
4
5
This script creates a pyproject.toml with the specified package as a dependency,
6
runs `uv run pywrangler sync` to download and prepare the package, and then
7
creates a zip file of the resulting python_modules directory.
8
"""
9
 
10
import argparse
11
import os
12
import shutil
13
import sys
14
import zipfile
15
from pathlib import Path
16
from typing import Literal
17
 
18
from tool_utils import hexdigest, run
19
 
20
PYPROJECT_TEMPLATE = """[project]
21
name = "vendor-test"
22
version = "0.1.0"
23
requires-python = ">=3.12"
24
dependencies = {dependencies}
25
26
[dependency-groups]
27
dev = ["workers-py>=1.6"]
28
"""
29
 
30
WRANGLER_TOML = """
31
name = "hello-python-bindings"
32
main = "src/entry.py"
33
compatibility_flags = {compat_flags}
34
compatibility_date = "2025-08-14"
35
"""
36
 
37
type PyVer = Literal["3.12", "3.13"]
38
 
39
 
40
def create_pyproject_toml(package_names: list[str], target_dir: Path) -> Path:
41
    """Create a pyproject.toml file with the specified package as a dependency."""
42
    pyproject_content = PYPROJECT_TEMPLATE.format(dependencies=repr(package_names))
43
    pyproject_path = target_dir / "pyproject.toml"
44
    pyproject_path.write_text(pyproject_content)
45
    return pyproject_path
46
 
47
 
48
def create_wrangler_toml(target_dir: Path, python: PyVer):
49
    compat_flags = ["python_workers"]
50
    match python:
51
        case "3.13":
52
            compat_flags.append("python_workers_20250116")
53
 
54
    (target_dir / "wrangler.toml").write_text(
55
        WRANGLER_TOML.format(compat_flags=repr(compat_flags))
56
    )
57
 
58
 
59
def run_pywrangler_sync(work_dir: Path) -> Path:
60
    """Run `uv run pywrangler sync` in the specified directory."""
61
    env = os.environ.copy()
62
    env["_PYODIDE_EXTRA_MOUNTS"] = str(work_dir)
63
    # TODO: Make pywrangler understand how to use Python 3.13 correctly and
64
    # remove these extra commands
65
    run(["uv", "run", "pywrangler", "sync"], cwd=work_dir, env=env)
66
    python_modules_dir = work_dir / "python_modules"
67
    if not python_modules_dir.exists():
68
        print(f"Error: python_modules directory not found at {python_modules_dir}")
69
        sys.exit(1)
70
    return python_modules_dir
71
 
72
 
73
def create_zip_archive(source_dir: Path, output_dir: Path) -> bool:
74
    """Create a zip archive of the python_modules directory.
75
76
    Return value indicates whether the archive includes any binary modules (.so
77
    files).
78
    """
79
    tmp_path = output_dir / "tmp.zip"
80
    native = False
81
 
82
    with zipfile.ZipFile(tmp_path, "w", zipfile.ZIP_DEFLATED) as zipf:
83
        for file_path in source_dir.rglob("*"):
84
            if file_path.is_file():
85
                # Store relative path from python_modules directory
86
                arcname = file_path.relative_to(source_dir)
87
                zipf.write(file_path, arcname)
88
                if file_path.suffix == ".so":
89
                    native = True
90
 
91
    return native
92
 
93
 
94
def vendor_package(package_names: list[str], python: PyVer) -> tuple[Path, bool]:
95
    """Main function to vendor a Python package."""
96
    tmp_dir = Path("/tmp")
97
    vendor_name = "-".join(package_names)
98
    work_dir = tmp_dir / f"vendor-{vendor_name}"
99
 
100
    # Clean up any existing work directory
101
    if work_dir.exists():
102
        shutil.rmtree(work_dir)
103
    work_dir.mkdir(parents=True)
104
 
105
    try:
106
        # Create pyproject.toml
107
        print(f"Creating pyproject.toml in {work_dir}")
108
        create_pyproject_toml(package_names, work_dir)
109
        create_wrangler_toml(work_dir, python)
110
 
111
        # Run pywrangler sync
112
        print("Running uv run pywrangler sync...")
113
        python_modules_dir = run_pywrangler_sync(work_dir)
114
 
115
        # Create zip archive
116
        print("Creating zip archive...")
117
        native = create_zip_archive(python_modules_dir, tmp_dir)
118
        py = f"-{python}" if native else ""
119
        zip_name = f"{vendor_name}{py}-vendored-for-ew-testing.zip"
120
        zip_path = tmp_dir / zip_name
121
        shutil.move(tmp_dir / "tmp.zip", zip_path)
122
    except Exception as e:
123
        print(f"Error vendoring packages {package_names!r}: {e}")
124
        sys.exit(1)
125
    else:
126
        print(f"Successfully created: {zip_path}")
127
        print(
128
            "Upload this zip file to the ew-snapshot-tests R2 bucket: "
129
            + "https://dash.cloudflare.com/e415f1017791ced9d5f3eb0df2b31c9e/r2/default/buckets/ew-snapshot-tests"
130
        )
131
        return zip_path, native
132
    finally:
133
        # Clean up work directory
134
        if work_dir.exists():
135
            shutil.rmtree(work_dir)
136
 
137
 
138
def main() -> int:
139
    """Main entry point."""
140
    parser = argparse.ArgumentParser(
141
        description="Create a zip file of a vendored Python package's source files for vendored_py_wd_test."
142
    )
143
    parser.add_argument(
144
        "package_name", help="Name of the Python package to vendor", nargs="*"
145
    )
146
    parser.add_argument("-p", "--python", help="Name of the Python version to use")
147
 
148
    args = parser.parse_args()
149
    if args.python is None:
150
        args.python = "3.12"
151
 
152
    if not args.package_name:
153
        print("Error: Package name is required")
154
        return 1
155
 
156
    try:
157
        zip_path, native = vendor_package(args.package_name, args.python)
158
        print("Update python_metadata.bzl with:\n")
159
        abi = args.python if native else None
160
        i1 = " " * 12
161
        i2 = " " * 16
162
        print(i1 + "{")
163
        print(i2 + f'"name": "{"-".join(args.package_name)}",')
164
        print(i2 + f'"abi": {abi!r},')
165
        print(i2 + f'"sha256": "{hexdigest(zip_path)}",')
166
        print(i1 + "},")
167
        print()
168
    except KeyboardInterrupt:
169
        print("\nOperation cancelled by user")
170
        return 1
171
    else:
172
        return 0
173
 
174
 
175
if __name__ == "__main__":
176
    sys.exit(main())