zhenxun_bot/zhenxun/utils/browser.py

import os
import sys

from nonebot import get_driver
from playwright.__main__ import main
from playwright.async_api import Browser, Playwright, async_playwright
from zhenxun.configs.config import SYSTEM_PROXY

from zhenxun.services.log import logger

driver = get_driver()

_playwright: Playwright | None = None
_browser: Browser | None = None


@driver.on_startup
async def start_browser():
    global _playwright
    global _browser
    install()
    await check_playwright_env()
    _playwright = await async_playwright().start()
    _browser = await _playwright.chromium.launch()


@driver.on_shutdown
async def shutdown_browser():
    if _browser:
        await _browser.close()
    if _playwright:
        await _playwright.stop()  # type: ignore


def get_browser() -> Browser:
    if not _browser:
        raise RuntimeError("playwright is not initalized")
    return _browser


def install():
    """自动安装、更新 Chromium"""

    def set_env_variables():
        os.environ[
            "PLAYWRIGHT_DOWNLOAD_HOST"] = "https://npmmirror.com/mirrors/playwright/"
        if SYSTEM_PROXY:
            os.environ["HTTPS_PROXY"] = SYSTEM_PROXY

    def restore_env_variables():
        os.environ.pop("PLAYWRIGHT_DOWNLOAD_HOST", None)
        if SYSTEM_PROXY:
            os.environ.pop("HTTPS_PROXY", None)
        if original_proxy is not None:
            os.environ["HTTPS_PROXY"] = original_proxy

    def try_install_chromium():
        try:
            sys.argv = ["", "install", "chromium"]
            main()
        except SystemExit as e:
            return e.code == 0
        return False

    logger.info("检查 Chromium 更新")

    original_proxy = os.environ.get("HTTPS_PROXY")
    set_env_variables()

    success = try_install_chromium()

    if not success:
        logger.info("Chromium 更新失败，尝试从原始仓库下载，速度较慢")
        os.environ["PLAYWRIGHT_DOWNLOAD_HOST"] = ""
        success = try_install_chromium()

    restore_env_variables()

    if not success:
        raise RuntimeError("未知错误，Chromium 下载失败")


async def check_playwright_env():
    """检查 Playwright 依赖"""
    logger.info("检查 Playwright 依赖")
    try:
        async with async_playwright() as p:
            await p.chromium.launch()
    except Exception as e:
        raise ImportError("加载失败，Playwright 依赖不全，") from e