Coverage for src/mlopus/utils/iter_utils.py: 97%

1import collections.abc 

2from typing import Generic, Callable, TypeVar, Iterable, Any, List 

3 

4T = TypeVar("T") 

5 

6R = TypeVar("R") # Type of result 

7Token = Any | None # Type of token 

8 

9 

10class Page(Generic[R]): 

11 """Page container to be used with Paginator.""" 

12 

13 def __init__(self, token: Token, results: List[R]): 

14 self.token = token 

15 self.results = results 

16 

17 def map_results(self, mapper: Callable[[R], T]) -> "Page": 

18 """Get same page with each result mapped by callback.""" 

19 return Page[T](self.token, [mapper(x) for x in self.results]) 

20 

21 def map_page(self, mapper: Callable[[R], T]) -> "Page": 

22 """Get same page with all results mapped by callback.""" 

23 return Page[T](self.token, mapper(self.results)) 

24 

25 

26class Paginator(Generic[R], collections.abc.Iterable[List[R]]): 

27 """Functional results paginator. 

28 

29 Example: 

30 page_size, max_results = 3, 9 

31 

32 def get_page(token): 

33 data = range(token := token or 0, new_token := token + page_size) 

34 return Page(new_token if new_token < max_results else None, data) 

35 

36 for x in Paginator(get_page).map_with(lambda x: x + 1): 

37 print(x) # 1, 2, 3, 4... 9 

38 """ 

39 

40 def __init__(self, get_page: Callable[[Token], Page[R]]): 

41 self._get_page = get_page 

42 

43 def __iter__(self) -> Iterable[List[str]]: # noqa 

44 token = None 

45 

46 while True: 

47 yield (page := self._get_page(token)).results 

48 

49 if not (token := page.token): 

50 break 

51 

52 def map_results(self, mapper: Callable[[R], T]) -> "Paginator": 

53 """Get same paginator with each result mapped by callback.""" 

54 return Paginator[T](lambda token: self._get_page(token).map_results(mapper)) 

55 

56 def map_pages(self, mapper: Callable[[List[R]], List[T]]) -> "Paginator": 

57 """Get same paginator with each page mapped by callback.""" 

58 return Paginator[T](lambda token: self._get_page(token).map_page(mapper)) 

59 

60 def collapse(self) -> "Paginator": 

61 """Get a paginator with all possible results in a single page.""" 

62 return Paginator[R](lambda _: Page(None, list(result for page in self for result in page))) 

63 

64 @classmethod 

65 def single_page(cls, page_data: Iterable[R]) -> "Paginator": 

66 """Get a paginator of a single page from a results iterable.""" 

67 return Paginator[R](lambda token: Page(token=None, results=list(page_data)))