| 1 | import json |
| 2 | import os |
| 3 | import re |
| 4 | import subprocess |
| 5 | import sys |
| 6 | from importlib import resources |
| 7 | from pathlib import Path |
| 8 | |
| 9 | from helpers import files |
| 10 | |
| 11 | FULL_CHROMIUM_PATTERNS = ( |
| 12 | "chromium-*/chrome-linux*/chrome", |
| 13 | "chromium-*/chrome-win*/chrome.exe", |
| 14 | ) |
| 15 | PLAYWRIGHT_CACHE_ENV = "A0_BROWSER_PLAYWRIGHT_CACHE_DIR" |
| 16 | PLAYWRIGHT_CACHE_DIR = ("tmp", "playwright") |
| 17 | RETIRED_PLAYWRIGHT_CACHE_DIRS = ( |
| 18 | ("usr", "plugins", "_browser", "playwright"), |
| 19 | ("usr", "browser", "playwright"), |
| 20 | ) |
| 21 | |
| 22 | |
| 23 | def _primary_cache_dir() -> Path: |
| 24 | override = os.environ.get(PLAYWRIGHT_CACHE_ENV, "").strip() |
| 25 | if override: |
| 26 | return Path(override).expanduser() |
| 27 | return Path(files.get_abs_path(*PLAYWRIGHT_CACHE_DIR)) |
| 28 | |
| 29 | |
| 30 | def get_playwright_cache_dir() -> str: |
| 31 | return str(_primary_cache_dir()) |
| 32 | |
| 33 | |
| 34 | def get_playwright_cache_dirs() -> list[Path]: |
| 35 | primary = _primary_cache_dir() |
| 36 | candidates = [primary, *get_retired_playwright_cache_dirs()] |
| 37 | seen: set[str] = set() |
| 38 | unique: list[Path] = [] |
| 39 | for candidate in candidates: |
| 40 | key = str(candidate) |
| 41 | if key in seen: |
| 42 | continue |
| 43 | seen.add(key) |
| 44 | unique.append(candidate) |
| 45 | return unique |
| 46 | |
| 47 | |
| 48 | def get_retired_playwright_cache_dirs() -> list[Path]: |
| 49 | return [Path(files.get_abs_path(*parts)) for parts in RETIRED_PLAYWRIGHT_CACHE_DIRS] |
| 50 | |
| 51 | |
| 52 | def configure_playwright_env() -> str: |
| 53 | cache_dir = get_playwright_cache_dir() |
| 54 | Path(cache_dir).mkdir(parents=True, exist_ok=True) |
| 55 | os.environ["PLAYWRIGHT_BROWSERS_PATH"] = cache_dir |
| 56 | return cache_dir |
| 57 | |
| 58 | |
| 59 | def find_playwright_binary(cache_dir: Path, revision: str = "") -> Path | None: |
| 60 | prefix = f"chromium-{revision}" if revision.isdigit() else "chromium-*" |
| 61 | binaries = [ |
| 62 | binary |
| 63 | for pattern in FULL_CHROMIUM_PATTERNS |
| 64 | for binary in cache_dir.glob(pattern.replace("chromium-*", prefix)) |
| 65 | if binary.exists() |
| 66 | ] |
| 67 | return max(binaries, key=_chromium_revision) if binaries else None |
| 68 | |
| 69 | |
| 70 | def _chromium_revision(binary: Path) -> int: |
| 71 | match = re.search(r"chromium-(\d+)", binary.as_posix()) |
| 72 | return int(match.group(1)) if match else -1 |
| 73 | |
| 74 | |
| 75 | def get_playwright_binary() -> Path | None: |
| 76 | cache_dir = _primary_cache_dir() |
| 77 | binary = find_playwright_binary(_primary_cache_dir()) |
| 78 | revision = get_playwright_chromium_revision() |
| 79 | if revision and (not binary or _chromium_revision(binary) != int(revision)): |
| 80 | return find_playwright_binary(cache_dir, revision=revision) |
| 81 | return binary |
| 82 | |
| 83 | |
| 84 | def get_playwright_chromium_revision() -> str: |
| 85 | try: |
| 86 | manifest = resources.files("patchright").joinpath("driver/package/browsers.json") |
| 87 | browsers = json.loads(manifest.read_text(encoding="utf-8"))["browsers"] |
| 88 | revision = next( |
| 89 | str(browser.get("revision", "")) |
| 90 | for browser in browsers |
| 91 | if browser.get("name") == "chromium" |
| 92 | ) |
| 93 | except (ImportError, FileNotFoundError, KeyError, StopIteration, TypeError, ValueError): |
| 94 | return "" |
| 95 | return revision if revision.isdigit() else "" |
| 96 | |
| 97 | |
| 98 | def ensure_playwright_binary() -> Path: |
| 99 | binary = get_playwright_binary() |
| 100 | if binary: |
| 101 | return binary |
| 102 | |
| 103 | cache_dir = configure_playwright_env() |
| 104 | env = os.environ.copy() |
| 105 | env["PLAYWRIGHT_BROWSERS_PATH"] = cache_dir |
| 106 | subprocess.check_call( |
| 107 | [sys.executable, "-m", "patchright", "install", "chromium", "--no-shell"], |
| 108 | env=env, |
| 109 | ) |
| 110 | |
| 111 | binary = get_playwright_binary() |
| 112 | if not binary: |
| 113 | raise RuntimeError("Patchright Chromium binary not found after installation") |
| 114 | return binary |