Coverage for human_requests / tools / base.py: 78%

1import inspect 

2from functools import wraps 

3from typing import Any, Awaitable, Callable, Type, TypeVar 

4 

5from playwright.async_api import Error as PlaywrightError 

6 

7F = TypeVar("F", bound=Callable[..., Awaitable[Any]]) 

8T = TypeVar("T", bound=Type) 

9 

10 

11def make_screenshot(method: F) -> F: 

12 @wraps(method) 

13 async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: 

14 try: 

15 return await method(self, *args, **kwargs) 

16 except PlaywrightError: 

17 try: 

18 from ..human_page import HumanPage 

19 

20 if not isinstance(self, HumanPage) or self.on_error_screenshot_path: 

21 await self.screenshot(path=self.on_error_screenshot_path or "error.png") 

22 except Exception as screenshot_error: 

23 print(f"Screenshot failed: {screenshot_error}") 

24 raise 

25 

26 return wrapper # type: ignore[return-value] 

27 

28 

29def auto_wrap_methods(decorator: Callable) -> Callable[[T], T]: 

30 """ 

31 Фабрика декораторов класса. Применяет декоратор только к асинхронным методам класса, 

32 кроме магических методов. Синхронные методы не оборачиваются. 

33 """ 

34 

35 def class_decorator(cls: T) -> T: 

36 for attr_name in dir(cls): 

37 # Пропускаем магические методы 

38 if attr_name.startswith("__") and attr_name.endswith("__"): 

39 continue 

40 attr = getattr(cls, attr_name) 

41 if not callable(attr): 

42 continue 

43 # Оборачиваем ТОЛЬКО асинхронные методы 

44 if not inspect.iscoroutinefunction(attr): 

45 continue 

46 # Если метод уже обёрнут (имеет __wrapped__), пропускаем 

47 if hasattr(attr, "__wrapped__"): 

48 continue 

49 wrapped = decorator(attr) 

50 setattr(cls, attr_name, wrapped) 

51 return cls 

52 

53 return class_decorator