ext_browser_context.py

import os
import logging
import threading
from typing import Optional, List, TYPE_CHECKING
from playwright.sync_api import Playwright, BrowserType, Browser, BrowserContext, Page
from ext_playwright import ExtPlaywrightOption, get_free_port, connect_over_cdp

if TYPE_CHECKING:
    from ext_page import ExtPage as ExtPageType

logger = logging.getLogger(__name__)


class ExtBrowserContext:
    """扩展浏览器上下文类"""

    def __init__(
        self,
        pw: Playwright,
        browser_type: BrowserType,
        browser: Optional[Browser],
        browser_context: BrowserContext
    ):
        self.pw = pw
        self.browser_type = browser_type
        self.browser = browser
        self.browser_context = browser_context
        self.ext_pages: List[ExtPage] = []
        self._lock = threading.Lock()

        # 初始化已有页面
        from ext_page import ExtPage
        for page in browser_context.pages:
            self.ext_pages.append(ExtPage(page, self, locked=False))

    @classmethod
    def new_debug_ext_browser_context(cls, ext_pw_opt: ExtPlaywrightOption) -> 'ExtBrowserContext':
        """创建调试模式的浏览器上下文

        Args:
            ext_pw_opt: 扩展选项

        Returns:
            ExtBrowserContext 实例
        """
        from playwright.sync_api import sync_playwright

        ext_pw_opt.skip_install_browsers = True

        with sync_playwright() as pw:
            if not ext_pw_opt.user_data_dir:
                user_home_dir = os.path.expanduser("~")
                ext_pw_opt.user_data_dir = os.path.join(user_home_dir, "ChromeProfile")

            if ext_pw_opt.remote_debugging_port > 0:
                remote_debugging_port = ext_pw_opt.remote_debugging_port
            else:
                remote_debugging_port = get_free_port()

            browser_context = pw.chromium.launch_persistent_context(
                ext_pw_opt.user_data_dir,
                args=[f"--remote-debugging-port={remote_debugging_port}"],
                executable_path=ext_pw_opt.must_get_browser_path(),
                headless=ext_pw_opt.headless
            )

            return cls.build_ext_browser_context(pw, pw.chromium, None, browser_context)

    @classmethod
    def connect_debug_ext_browser_context(cls, ext_pw_opt: ExtPlaywrightOption) -> 'ExtBrowserContext':
        """连接到调试模式的浏览器

        Args:
            ext_pw_opt: 扩展选项

        Returns:
            ExtBrowserContext 实例
        """
        from playwright.sync_api import sync_playwright

        with sync_playwright() as pw:
            browser = connect_over_cdp(pw, ext_pw_opt)

            if browser is None:
                raise RuntimeError("browser connection is nil")

            if not browser.is_connected():
                browser.close()
                raise RuntimeError("browser is not connected")

            # 获取浏览器上下文
            contexts = browser.contexts
            if len(contexts) > 0:
                browser_context = contexts[0]
            else:
                browser_context = browser.new_context()

            ext_browser_context = cls.build_ext_browser_context(
                pw, pw.chromium, browser, browser_context
            )

            return ext_browser_context

    @classmethod
    def new_ext_browser_context(cls, ext_pw_opt: ExtPlaywrightOption) -> 'ExtBrowserContext':
        """创建新的浏览器上下文

        Args:
            ext_pw_opt: 扩展选项

        Returns:
            ExtBrowserContext 实例
        """
        from playwright.sync_api import sync_playwright

        with sync_playwright() as pw:
            browser_type = ext_pw_opt.get_browser_type(pw)

            executable_path = ext_pw_opt.browser_path if ext_pw_opt.browser_path else None
            channel = ext_pw_opt.channel if ext_pw_opt.channel else None

            browser_context = None
            browser = None

            if ext_pw_opt.user_data_dir:
                browser_context = browser_type.launch_persistent_context(
                    ext_pw_opt.user_data_dir,
                    executable_path=executable_path,
                    channel=channel,
                    headless=ext_pw_opt.headless
                )
            else:
                browser = browser_type.launch(
                    executable_path=executable_path,
                    channel=channel,
                    headless=ext_pw_opt.headless
                )

            return cls.build_ext_browser_context(pw, browser_type, browser, browser_context)

    @classmethod
    def build_ext_browser_context(
        cls,
        pw: Playwright,
        browser_type: BrowserType,
        browser: Optional[Browser],
        browser_context: Optional[BrowserContext]
    ) -> 'ExtBrowserContext':
        """构建扩展浏览器上下文

        Args:
            pw: Playwright 实例
            browser_type: 浏览器类型
            browser: 浏览器实例
            browser_context: 浏览器上下文实例

        Returns:
            ExtBrowserContext 实例
        """
        if pw is None:
            raise RuntimeError("playwright instance is nil")

        if browser_type is None:
            browser_type = pw.chromium

        if browser_context is None:
            if browser is None:
                browser = browser_type.launch()

            if not browser.is_connected():
                raise RuntimeError("browser is not connected")

            contexts = browser.contexts
            if len(contexts) > 0:
                browser_context = contexts[0]
            else:
                browser_context = browser.new_context()

        if browser_context is None:
            raise RuntimeError("browser context is nil")

        ext_bc = cls(pw, browser_type, browser, browser_context)

        # 初始化已有页面
        for page in browser_context.pages:
            ext_bc.ext_pages.append(ExtPage(page, ext_bc, locked=False))

        return ext_bc

    def close(self):
        """关闭浏览器上下文"""
        self.browser_context.close()
        if self.browser:
            self.browser.close()
        self.pw.stop()

    def get_or_new_ext_page(self) -> 'ExtPageType':
        """获取或创建新的扩展页面

        Returns:
            ExtPage 实例
        """
        # 清理已关闭的页面
        opened_pages = []
        for ext_page in self.ext_pages:
            if ext_page.is_closed:
                ext_page.close()
                ext_page.locked = False
            else:
                opened_pages.append(ext_page)

        self.ext_pages = opened_pages

        # 尝试获取未锁定的页面
        if len(self.ext_pages) > 0:
            with self._lock:
                for ext_page in self.ext_pages:
                    if not ext_page.locked:
                        ext_page.locked = True
                        return ext_page

        # 创建新页面
        return self.new_ext_page()

    def new_ext_page(self) -> 'ExtPageType':
        """创建新的扩展页面

        Returns:
            ExtPage 实例
        """
        page = self.browser_context.new_page()
        return self.build_ext_page(page)

    def build_ext_page(self, page: Page) -> 'ExtPage':
        """构建扩展页面

        Args:
            page: Playwright Page 实例

        Returns:
            ExtPage 实例
        """
        from ext_page import ExtPage
        ext_page = ExtPage(page, self, locked=True)
        self.ext_pages.append(ext_page)
        return ext_page

    def __getattr__(self, name):
        """代理到 Playwright BrowserContext"""
        return getattr(self.browser_context, name)