crawlab_feader/FworkSpider/feapder/network/downloader/_drissionpage.py

# -*- coding: utf-8 -*-
"""
Created on 2025-05-14 
---------
@summary:
---------
@author: Dzr
"""
import feapder.setting as setting
import feapder.utils.tools as tools
from feapder.network.downloader.base import RenderDownloader
from feapder.network.response import Response
from feapder.utils.webdriver import WebDriverPool, DrissionPageDriver


class DrissionPageDownloader(RenderDownloader):
    webdriver_pool: WebDriverPool = None

    @property
    def _webdriver_pool(self):
        if not self.__class__.webdriver_pool:
            self.__class__.webdriver_pool = WebDriverPool(
                **setting.DRISSIONPAGE, driver_cls=DrissionPageDriver
            )

        return self.__class__.webdriver_pool

    def download(self, request) -> Response:
        # 代理优先级 自定义 > 配置文件 > 随机
        if request.custom_proxies:
            proxy = request.get_proxy()
        elif setting.DRISSIONPAGE.get("proxy"):
            proxy = setting.DRISSIONPAGE.get("proxy")
        else:
            proxy = request.get_proxy()

        # user_agent优先级 自定义 > 配置文件 > 随机
        if request.custom_ua:
            user_agent = request.get_user_agent()
        elif setting.DRISSIONPAGE.get("user_agent"):
            user_agent = setting.DRISSIONPAGE.get("user_agent")
        else:
            user_agent = request.get_user_agent()

        cookies = request.get_cookies()
        url = request.url
        render_time = request.render_time or setting.DRISSIONPAGE.get("render_time")
        if request.get_params():
            url = tools.joint_url(url, request.get_params())

        driver: DrissionPageDriver = self._webdriver_pool.get(
            user_agent=user_agent, proxy=proxy
        )

        try:
            if cookies:
                driver.url = url
                driver.tab.set.cookies = cookies

            retry = setting.DRISSIONPAGE.get("retry")
            interval = setting.DRISSIONPAGE.get("interval")
            timeout = setting.DRISSIONPAGE.get("page_load")
            driver.tab.get(url, retry=retry, interval=interval, timeout=timeout)
            driver.wait_for_dom_load(timeout=render_time)

            if driver.tab.mode == "s":
                response = Response(driver.tab.response)
            else:
                cookies = driver.tab.cookies().as_dict()
                html = driver.tab.html
                response = Response.from_dict(
                    {
                        "url": driver.tab.url,
                        "cookies": cookies,
                        "_content": html.encode(),
                        "status_code": 200,
                        "elapsed": 666,
                        "headers": {
                            "User-Agent": driver.tab.user_agent,
                            "Cookie": tools.cookies2str(cookies),
                        },
                    }
                )

            response.driver = driver
            response.browser = driver
            return response
        except Exception as e:
            self._webdriver_pool.remove(driver)
            request.del_proxy()
            raise e

    def close(self, driver):
        if driver:
            self._webdriver_pool.remove(driver)

    def put_back(self, driver):
        """
        释放浏览器对象
        """
        self._webdriver_pool.put(driver)

    def close_all(self):
        """
        关闭所有浏览器
        """
        self._webdriver_pool.close()