Coverage for src/lilbee/crawler/backends/crawlberg/fetcher.py: 100%

1"""crawlberg-backed implementation of :class:`lilbee.crawler.fetcher.WebFetcher`.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import logging 

7from collections.abc import AsyncGenerator, Sequence 

8from contextlib import aclosing 

9from typing import TYPE_CHECKING, Any 

10 

11from lilbee.core.config.enums import CrawlRenderMode 

12from lilbee.crawler import url_filter 

13from lilbee.crawler.backends.crawlberg.config import _crawl_config, _CrawlSpec 

14from lilbee.crawler.backends.crawlberg.events import ( 

15 _SEED_DEPTH, 

16 _Event, 

17 _EventKind, 

18 _events, 

19) 

20from lilbee.crawler.backends.crawlberg.links import _ExcludedLinks 

21from lilbee.crawler.models import ( 

22 NO_LOGIN, 

23 CancelToken, 

24 ConcurrencySpec, 

25 CrawlLogin, 

26 FetchedPage, 

27 FilterSpec, 

28) 

29 

30if TYPE_CHECKING: 

31 from lilbee.crawler.fetcher import WebFetcher 

32 

33log = logging.getLogger(__name__) 

34 

35 

36_NO_CONTENT = "No content extracted" 

37 

38 

39def _url_allowed(url: str) -> bool: 

40 """True when lilbee's own URL policy admits *url*.""" 

41 try: 

42 url_filter.validate_crawl_url(url) 

43 except ValueError: 

44 return False 

45 return True 

46 

47 

48async def _to_page(event: _Event, seed_url: str) -> FetchedPage | None: 

49 """The fetched page for *event*, or None when lilbee's URL policy refuses its URL.""" 

50 if event.kind is _EventKind.ERROR: 

51 return FetchedPage(url=event.url, success=False, error=event.error) 

52 if not await asyncio.to_thread(_url_allowed, event.url): 

53 log.warning("Dropped %s: lilbee's URL policy refuses it", event.url) 

54 return None 

55 url = seed_url if event.depth == _SEED_DEPTH else event.url 

56 return FetchedPage(url=url, markdown=event.markdown, title=event.title) 

57 

58 

59def _single_result(page: FetchedPage | None, url: str) -> FetchedPage: 

60 """A single-URL fetch result under the requested URL, failed when it has no text.""" 

61 if page is not None and not page.success: 

62 return FetchedPage(url=url, success=False, error=page.error) 

63 if page is not None and page.markdown.strip(): 

64 return FetchedPage(url=url, markdown=page.markdown.strip(), title=page.title) 

65 return FetchedPage(url=url, success=False, error=_NO_CONTENT) 

66 

67 

68class CrawlbergFetcher: 

69 """:class:`WebFetcher` implementation backed by crawlberg.""" 

70 

71 def __init__( 

72 self, 

73 *, 

74 render_mode: CrawlRenderMode, 

75 chrome_args: Sequence[str] = (), 

76 user_agent: str = "", 

77 respect_robots_txt: bool = True, 

78 extra_wait_ms: int = 0, 

79 login: CrawlLogin = NO_LOGIN, 

80 ) -> None: 

81 self._render_mode = render_mode 

82 self._chrome_args = tuple(chrome_args) 

83 self._user_agent = user_agent 

84 self._respect_robots_txt = respect_robots_txt 

85 self._extra_wait_ms = extra_wait_ms 

86 self._login = login 

87 

88 async def __aenter__(self) -> CrawlbergFetcher: 

89 return self 

90 

91 async def __aexit__( 

92 self, 

93 exc_type: type[BaseException] | None, 

94 exc: BaseException | None, 

95 tb: Any, 

96 ) -> None: 

97 return None 

98 

99 def _stream( 

100 self, spec: _CrawlSpec, seed_url: str, cancel: CancelToken | None 

101 ) -> AsyncGenerator[_Event, None]: 

102 """This crawl's events; browser mode finds the headless shell before any engine starts.""" 

103 config = _crawl_config( 

104 self._render_mode, 

105 spec, 

106 self._chrome_args, 

107 self._user_agent, 

108 self._respect_robots_txt, 

109 self._extra_wait_ms, 

110 self._login, 

111 ) 

112 return _events(config, seed_url, cancel) 

113 

114 async def fetch_single( 

115 self, url: str, *, timeout: float, concurrency: ConcurrencySpec 

116 ) -> FetchedPage: 

117 """Fetch *url* alone; a redirected page keeps the requested URL, without its userinfo.""" 

118 spec = _CrawlSpec(depth=_SEED_DEPTH, max_pages=1, timeout=timeout, concurrency=concurrency) 

119 shown = url_filter.without_userinfo(url) 

120 async with aclosing(self._stream(spec, url, None)) as events: 

121 async for event in events: 

122 return _single_result(await _to_page(event, shown), shown) 

123 return _single_result(None, shown) 

124 

125 async def fetch_recursive( 

126 self, 

127 seed_url: str, 

128 *, 

129 depth: int | None, 

130 max_pages: int | None, 

131 timeout: float, 

132 concurrency: ConcurrencySpec, 

133 filters: FilterSpec, 

134 cancel: CancelToken | None = None, 

135 ) -> AsyncGenerator[FetchedPage, None]: 

136 """Stream pages crawlberg reaches breadth-first from *seed_url*. 

137 

138 A discovered link crawlberg's SSRF policy refuses is dropped without a page, 

139 like a link filtered out before it is fetched; a refused seed is a failed page. 

140 """ 

141 spec = _CrawlSpec(depth, max_pages, timeout, concurrency, filters) 

142 shown = url_filter.without_userinfo(seed_url) 

143 excluded = _ExcludedLinks(spec, shown) 

144 if excluded.record_seed(shown): 

145 return 

146 try: 

147 async with aclosing(self._stream(spec, seed_url, cancel)) as events: 

148 async for event in events: 

149 excluded.record(event) 

150 if event.refused_by_ssrf and event.url != shown: 

151 log.debug("crawlberg refused %s: %s", event.url, event.error) 

152 continue 

153 page = await _to_page(event, shown) 

154 if page is not None: 

155 yield page 

156 finally: 

157 excluded.summarize() 

158 

159 

160# Protocol conformance check: CrawlbergFetcher is structurally a WebFetcher. 

161if TYPE_CHECKING: 

162 _: WebFetcher = CrawlbergFetcher(render_mode=CrawlRenderMode.HTTP)