Coverage for notion_client/helpers.py: 100%

1"""Utility functions for notion-sdk-py.""" 

2from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, Generator, List 

3from urllib.parse import urlparse 

4from uuid import UUID 

5 

6 

7def pick(base: Dict[Any, Any], *keys: str) -> Dict[Any, Any]: 

8 """Return a dict composed of key value pairs for keys passed as args.""" 

9 result = {} 

10 for key in keys: 

11 if key not in base: 

12 continue 

13 value = base.get(key) 

14 if value is None and key == "start_cursor": 

15 continue 

16 result[key] = value 

17 return result 

18 

19 

20def get_url(object_id: str) -> str: 

21 """Return the URL for the object with the given id.""" 

22 return f"https://notion.so/{UUID(object_id).hex}" 

23 

24 

25def get_id(url: str) -> str: 

26 """Return the id of the object behind the given URL.""" 

27 parsed = urlparse(url) 

28 if parsed.netloc not in ("notion.so", "www.notion.so"): 

29 raise ValueError("Not a valid Notion URL.") 

30 path = parsed.path 

31 if len(path) < 32: 

32 raise ValueError("The path in the URL seems to be incorrect.") 

33 raw_id = path[-32:] 

34 return str(UUID(raw_id)) 

35 

36 

37def iterate_paginated_api( 

38 function: Callable[..., Any], **kwargs: Any 

39) -> Generator[List[Any], None, None]: 

40 """Return an iterator over the results of any paginated Notion API.""" 

41 next_cursor = kwargs.pop("start_cursor", None) 

42 

43 while True: 

44 response = function(**kwargs, start_cursor=next_cursor) 

45 for result in response.get("results"): 

46 yield result 

47 

48 next_cursor = response.get("next_cursor") 

49 if not response.get("has_more") or not next_cursor: 

50 return 

51 

52 

53def collect_paginated_api(function: Callable[..., Any], **kwargs: Any) -> List[Any]: 

54 """Collect all the results of paginating an API into a list.""" 

55 return [result for result in iterate_paginated_api(function, **kwargs)] 

56 

57 

58async def async_iterate_paginated_api( 

59 function: Callable[..., Awaitable[Any]], **kwargs: Any 

60) -> AsyncGenerator[List[Any], None]: 

61 """Return an async iterator over the results of any paginated Notion API.""" 

62 next_cursor = kwargs.pop("start_cursor", None) 

63 

64 while True: 

65 response = await function(**kwargs, start_cursor=next_cursor) 

66 for result in response.get("results"): 

67 yield result 

68 

69 next_cursor = response.get("next_cursor") 

70 if (not response["has_more"]) | (next_cursor is None): 

71 return 

72 

73 

74async def async_collect_paginated_api( 

75 function: Callable[..., Awaitable[Any]], **kwargs: Any 

76) -> List[Any]: 

77 """Collect asynchronously all the results of paginating an API into a list.""" 

78 return [result async for result in async_iterate_paginated_api(function, **kwargs)] 

79 

80 

81def is_full_block(response: Dict[Any, Any]) -> bool: 

82 """Return `True` if response is a full block.""" 

83 return response.get("object") == "block" and "type" in response 

84 

85 

86def is_full_page(response: Dict[Any, Any]) -> bool: 

87 """Return `True` if response is a full page.""" 

88 return response.get("object") == "page" and "url" in response 

89 

90 

91def is_full_database(response: Dict[Any, Any]) -> bool: 

92 """Return `True` if response is a full database.""" 

93 return response.get("object") == "database" and "title" in response 

94 

95 

96def is_full_page_or_database(response: Dict[Any, Any]) -> bool: 

97 """Return `True` if `response` is a full database or a full page.""" 

98 if response.get("object") == "database": 

99 return is_full_database(response) 

100 return is_full_page(response) 

101 

102 

103def is_full_user(response: Dict[Any, Any]) -> bool: 

104 """Return `True` if response is a full user.""" 

105 return "type" in response 

106 

107 

108def is_full_comment(response: Dict[Any, Any]) -> bool: 

109 """Return `True` if response is a full comment.""" 

110 return "type" in response 

111 

112 

113def is_text_rich_text_item_response(rich_text: Dict[Any, Any]) -> bool: 

114 """Return `True` if `rich_text` is a text.""" 

115 return rich_text.get("type") == "text" 

116 

117 

118def is_equation_rich_text_item_response(rich_text: Dict[Any, Any]) -> bool: 

119 """Return `True` if `rich_text` is an equation.""" 

120 return rich_text.get("type") == "equation" 

121 

122 

123def is_mention_rich_text_item_response(rich_text: Dict[Any, Any]) -> bool: 

124 """Return `True` if `rich_text` is a mention.""" 

125 return rich_text.get("type") == "mention"