From b3ad70de8200cc7c787953b69c6f612c5b6c6d37 Mon Sep 17 00:00:00 2001 From: Jakob de Maeyer Date: Thu, 12 Sep 2024 10:34:44 +0200 Subject: [PATCH] Allow custom PageMethod callbacks --- scrapy_playwright/handler.py | 6 +++++- scrapy_playwright/page.py | 6 +++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/scrapy_playwright/handler.py b/scrapy_playwright/handler.py index 0363f2b2..0e6ae7f9 100644 --- a/scrapy_playwright/handler.py +++ b/scrapy_playwright/handler.py @@ -5,6 +5,7 @@ import warnings from contextlib import suppress from dataclasses import dataclass, field as dataclass_field +from functools import partial from ipaddress import ip_address from time import time from typing import Awaitable, Callable, Dict, Optional, Tuple, Type, TypeVar, Union @@ -607,7 +608,10 @@ async def _apply_page_methods(self, page: Page, request: Request, spider: Spider for pm in page_methods: if isinstance(pm, PageMethod): try: - method = getattr(page, pm.method) + if callable(pm.method): + method = partial(pm.method, page) + else: + method = getattr(page, pm.method) except AttributeError as ex: logger.warning( "Ignoring %r: could not find method", diff --git a/scrapy_playwright/page.py b/scrapy_playwright/page.py index 0d016a15..8112bf1b 100644 --- a/scrapy_playwright/page.py +++ b/scrapy_playwright/page.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, Callable __all__ = ["PageMethod"] @@ -10,8 +10,8 @@ class PageMethod: Playwright page, such as "click", "screenshot", "evaluate", etc. """ - def __init__(self, method: str, *args, **kwargs) -> None: - self.method: str = method + def __init__(self, method: str | Callable, *args, **kwargs) -> None: + self.method: str | Callable = method self.args: tuple = args self.kwargs: dict = kwargs self.result: Any = None