Coverage for notion_client/helpers.py: 100%

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

2 

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

4from urllib.parse import urlparse 

5from uuid import UUID 

6 

7 

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

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

10 result = {} 

11 for key in keys: 

12 if key not in base: 

13 continue 

14 value = base.get(key) 

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

16 continue 

17 result[key] = value 

18 return result 

19 

20 

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

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

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

24 

25 

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

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

28 parsed = urlparse(url) 

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

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

31 path = parsed.path 

32 if len(path) < 32: 

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

34 raw_id = path[-32:] 

35 return str(UUID(raw_id)) 

36 

37 

38def iterate_paginated_api( 

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

40) -> Generator[Any, None, None]: 

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

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

43 

44 while True: 

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

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

47 yield result 

48 

49 next_cursor = response.get("next_cursor") 

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

51 return 

52 

53 

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

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

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

57 

58 

59async def async_iterate_paginated_api( 

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

61) -> AsyncGenerator[Any, None]: 

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

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

64 

65 while True: 

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

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

68 yield result 

69 

70 next_cursor = response.get("next_cursor") 

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

72 return 

73 

74 

75async def async_collect_paginated_api( 

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

77) -> List[Any]: 

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

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

80 

81 

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

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

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

85 

86 

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

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

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

90 

91 

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

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

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

95 

96 

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

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

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

100 return is_full_database(response) 

101 return is_full_page(response) 

102 

103 

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

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

106 return "type" in response 

107 

108 

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

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

111 return "type" in response 

112 

113 

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

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

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

117 

118 

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

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

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

122 

123 

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

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

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