Coverage for human_requests/tools/helper_tools.py: 62%

1""" 

2helper_tools — helper utilities independent of a specific Session: 

3- assembling/merging storage_state (cookies + localStorage) 

4- unified navigation handler with soft retries 

5""" 

6 

7from __future__ import annotations 

8 

9from typing import TYPE_CHECKING, Awaitable, Callable, Literal, Optional 

10 

11from playwright.async_api import BrowserContext, Page 

12from playwright.async_api import TimeoutError as PlaywrightTimeoutError 

13 

14if TYPE_CHECKING: 

15 from playwright._impl._api_structures import LocalStorageEntry, OriginState 

16 from playwright.async_api import StorageState, StorageStateCookie 

17 

18 from ..abstraction.cookies import CookieManager 

19 

20# Зависящие типы простые и стабильные — импортируем прямо. 

21# CookieManager нужен только как протокол поведения (to_playwright/add_from_playwright). 

22 

23 

24def build_storage_state_for_context( 

25 *, 

26 local_storage: dict[str, dict[str, str]], 

27 cookie_manager: "CookieManager", 

28) -> "StorageState": 

29 cookie_list: list["StorageStateCookie"] = cookie_manager.to_playwright() 

30 origins: list["OriginState"] = [] 

31 

32 for origin, kv in local_storage.items(): 

33 if not kv: 

34 continue 

35 entries: list["LocalStorageEntry"] = [{"name": k, "value": v} for k, v in kv.items()] 

36 origins.append({"origin": origin, "localStorage": entries}) 

37 

38 return {"cookies": cookie_list, "origins": origins} 

39 

40 

41async def merge_storage_state_from_context( 

42 ctx: BrowserContext, *, cookie_manager: "CookieManager" 

43) -> dict[str, dict[str, str]]: 

44 """ 

45 Reads storage_state from the context and synchronizes internal state: 

46 - localStorage: FULL overwrite and returned outward 

47 - cookies: ADD/UPDATE in the provided CookieManager 

48 """ 

49 state = await ctx.storage_state() # dict с 'cookies' и 'origins' 

50 

51 # localStorage — точная перезапись 

52 new_ls: dict[str, dict[str, str]] = {} 

53 for o in state.get("origins", []) or []: 

54 origin = str(o.get("origin", "")) 

55 if not origin: 

56 continue 

57 kv: dict[str, str] = {} 

58 for pair in o.get("localStorage", []) or []: 

59 name = str(pair.get("name", "")) 

60 value = "" if pair.get("value") is None else str(pair.get("value")) 

61 if name: 

62 kv[name] = value 

63 new_ls[origin] = kv 

64 

65 # cookies — пополняем CookieManager 

66 cookies_list = state.get("cookies", []) or [] 

67 if cookies_list: 

68 cookie_manager.add_from_playwright(cookies_list) 

69 

70 return new_ls 

71 

72 

73async def handle_nav_with_retries( 

74 page: Page, 

75 *, 

76 target_url: str, 

77 wait_until: Literal["commit", "load", "domcontentloaded", "networkidle"], 

78 timeout_ms: int, 

79 attempts: int, 

80 on_retry: Optional[Callable[[], Awaitable[None]]] = None, 

81) -> None: 

82 """ 

83 Unified navigation handler with soft retries for goto/render. 

84 Catches ONLY PlaywrightTimeoutError. On retries calls on_retry() 

85 (if provided), then performs a reload (soft refresh). 

86 """ 

87 try: 

88 await page.goto(target_url, wait_until=wait_until, timeout=timeout_ms) 

89 except PlaywrightTimeoutError as last_err: 

90 while attempts > 0: 

91 attempts -= 1 

92 if on_retry is not None: 

93 await on_retry() 

94 try: 

95 await page.reload(wait_until=wait_until, timeout=timeout_ms) 

96 last_err = None # type: ignore[assignment] 

97 break 

98 except PlaywrightTimeoutError as e: 

99 last_err = e 

100 if last_err is not None: 

101 raise last_err