|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import hashlib |
| 4 | +import importlib.metadata |
| 5 | +import importlib.util |
| 6 | +import json |
| 7 | +import os |
| 8 | +import shutil |
| 9 | +import sys |
| 10 | +import tempfile |
| 11 | +import urllib.request |
| 12 | +import zipfile |
| 13 | +from pathlib import Path |
| 14 | + |
| 15 | +from packaging.tags import sys_tags |
| 16 | +from packaging.utils import canonicalize_name, parse_wheel_filename |
| 17 | + |
| 18 | +PACKAGE_NAME = "a3s-code" |
| 19 | +DIST_NAME = "a3s_code" |
| 20 | +REPO = "A3S-Lab/Code" |
| 21 | +MANIFEST_NAME = "python-native-manifest.json" |
| 22 | +ENV_BASE_URL = "A3S_CODE_RELEASE_BASE_URL" |
| 23 | +ENV_CACHE_DIR = "A3S_CODE_CACHE_DIR" |
| 24 | +ENV_OFFLINE = "A3S_CODE_OFFLINE" |
| 25 | + |
| 26 | + |
| 27 | +def _package_version() -> str: |
| 28 | + return importlib.metadata.version(PACKAGE_NAME) |
| 29 | + |
| 30 | + |
| 31 | +def _release_tag() -> str: |
| 32 | + return f"v{_package_version()}" |
| 33 | + |
| 34 | + |
| 35 | +def _release_base_url() -> str: |
| 36 | + override = os.environ.get(ENV_BASE_URL) |
| 37 | + if override: |
| 38 | + return override.rstrip("/") |
| 39 | + return f"https://github.com/{REPO}/releases/download/{_release_tag()}" |
| 40 | + |
| 41 | + |
| 42 | +def _cache_root() -> Path: |
| 43 | + override = os.environ.get(ENV_CACHE_DIR) |
| 44 | + if override: |
| 45 | + return Path(override).expanduser() |
| 46 | + if sys.platform == "win32": |
| 47 | + base = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData/Local")) |
| 48 | + elif sys.platform == "darwin": |
| 49 | + base = Path.home() / "Library/Caches" |
| 50 | + else: |
| 51 | + base = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) |
| 52 | + return base / "a3s-code" |
| 53 | + |
| 54 | + |
| 55 | +def _download_bytes(url: str) -> bytes: |
| 56 | + request = urllib.request.Request(url, headers={"User-Agent": "a3s-code-bootstrap/1"}) |
| 57 | + with urllib.request.urlopen(request) as response: |
| 58 | + return response.read() |
| 59 | + |
| 60 | + |
| 61 | +def _load_manifest() -> dict: |
| 62 | + manifest_url = f"{_release_base_url()}/{MANIFEST_NAME}" |
| 63 | + return json.loads(_download_bytes(manifest_url).decode("utf-8")) |
| 64 | + |
| 65 | + |
| 66 | +def _wheel_assets(manifest: dict) -> list[dict]: |
| 67 | + return list(manifest.get("assets", [])) |
| 68 | + |
| 69 | + |
| 70 | +def _select_asset(assets: list[dict]) -> dict: |
| 71 | + tag_list = list(sys_tags()) |
| 72 | + for tag in tag_list: |
| 73 | + for asset in assets: |
| 74 | + filename = asset["filename"] |
| 75 | + try: |
| 76 | + name, _, _, wheel_tags = parse_wheel_filename(filename) |
| 77 | + except Exception: |
| 78 | + continue |
| 79 | + if canonicalize_name(name) != canonicalize_name(DIST_NAME): |
| 80 | + continue |
| 81 | + if tag in wheel_tags: |
| 82 | + return asset |
| 83 | + raise ImportError("No compatible native a3s-code wheel found for this platform") |
| 84 | + |
| 85 | + |
| 86 | +def _sha256(path: Path) -> str: |
| 87 | + digest = hashlib.sha256() |
| 88 | + with path.open("rb") as file: |
| 89 | + for chunk in iter(lambda: file.read(1024 * 1024), b""): |
| 90 | + digest.update(chunk) |
| 91 | + return digest.hexdigest() |
| 92 | + |
| 93 | + |
| 94 | +def _download_wheel(asset: dict, destination: Path) -> Path: |
| 95 | + destination.mkdir(parents=True, exist_ok=True) |
| 96 | + wheel_path = destination / asset["filename"] |
| 97 | + if wheel_path.exists(): |
| 98 | + if _sha256(wheel_path) == asset["sha256"]: |
| 99 | + return wheel_path |
| 100 | + wheel_path.unlink() |
| 101 | + |
| 102 | + if os.environ.get(ENV_OFFLINE) == "1": |
| 103 | + raise ImportError(f"Offline mode enabled and native wheel is missing: {wheel_path}") |
| 104 | + |
| 105 | + tmp_dir = Path(tempfile.mkdtemp(prefix="a3s-code-wheel-")) |
| 106 | + try: |
| 107 | + tmp_path = tmp_dir / asset["filename"] |
| 108 | + tmp_path.write_bytes(_download_bytes(f"{_release_base_url()}/{asset['filename']}")) |
| 109 | + if _sha256(tmp_path) != asset["sha256"]: |
| 110 | + raise ImportError(f"SHA256 mismatch for {asset['filename']}") |
| 111 | + tmp_path.replace(wheel_path) |
| 112 | + finally: |
| 113 | + shutil.rmtree(tmp_dir, ignore_errors=True) |
| 114 | + return wheel_path |
| 115 | + |
| 116 | + |
| 117 | +def _extract_native_module(wheel_path: Path, destination: Path) -> Path: |
| 118 | + destination.mkdir(parents=True, exist_ok=True) |
| 119 | + with zipfile.ZipFile(wheel_path) as archive: |
| 120 | + native_members = [ |
| 121 | + member |
| 122 | + for member in archive.namelist() |
| 123 | + if member.startswith("a3s_code/") |
| 124 | + and "/_native" in member |
| 125 | + and member.endswith((".so", ".pyd")) |
| 126 | + ] |
| 127 | + if not native_members: |
| 128 | + raise ImportError(f"No native module found in wheel: {wheel_path.name}") |
| 129 | + member = native_members[0] |
| 130 | + extracted = destination / Path(member).name |
| 131 | + if not extracted.exists(): |
| 132 | + with archive.open(member) as src, extracted.open("wb") as dst: |
| 133 | + shutil.copyfileobj(src, dst) |
| 134 | + return extracted |
| 135 | + |
| 136 | + |
| 137 | +def load_native_module(): |
| 138 | + version = _package_version() |
| 139 | + cache_dir = _cache_root() / version |
| 140 | + assets = _wheel_assets(_load_manifest()) |
| 141 | + asset = _select_asset(assets) |
| 142 | + wheel_path = _download_wheel(asset, cache_dir / "wheels") |
| 143 | + native_path = _extract_native_module(wheel_path, cache_dir / "native") |
| 144 | + |
| 145 | + spec = importlib.util.spec_from_file_location("a3s_code._native", native_path) |
| 146 | + if spec is None or spec.loader is None: |
| 147 | + raise ImportError(f"Unable to load native module from {native_path}") |
| 148 | + module = importlib.util.module_from_spec(spec) |
| 149 | + sys.modules["a3s_code._native"] = module |
| 150 | + spec.loader.exec_module(module) |
| 151 | + return module |
0 commit comments