From f4c12393be87c210186d70af8e0220990b9f0893 Mon Sep 17 00:00:00 2001 From: Shannon Shen <22512825+lolipopshock@users.noreply.github.com> Date: Sun, 13 Feb 2022 20:10:47 -0500 Subject: [PATCH 1/2] Better RichTextObject --- src/notion_df/base.py | 117 ++++++++++++++++++++++++---------------- src/notion_df/values.py | 14 ++--- 2 files changed, 77 insertions(+), 54 deletions(-) diff --git a/src/notion_df/base.py b/src/notion_df/base.py index 5cb22de..eacb28a 100644 --- a/src/notion_df/base.py +++ b/src/notion_df/base.py @@ -50,53 +50,6 @@ class RichTextTypeEnum(str, Enum): Equation = "equation" -class Annotation(BaseModel): - bold: bool - italic: bool - strikethrough: bool - underline: bool - code: bool - color: NotionExtendedColorEnum - - -class BaseRichText(BaseModel): - plain_text: Optional[str] - # TODO: The Optional[plain_text] is used when creating property values - href: Optional[str] = None - annotations: Optional[Annotation] = None - type: Optional[RichTextTypeEnum] - - -class Link(BaseModel): - type: Optional[str] = "url" - url: str - - -class Text(BaseModel): - content: str - link: Optional[Link] - - -class RichText(BaseRichText): - text: Text - - @classmethod - def from_value(cls, value: str): - return cls(text=Text(content=value)) - - @classmethod - def encode_string(cls, value: str) -> List["RichText"]: - chunk_size = RICH_TEXT_CONTENT_MAX_LENGTH - return [ - cls(text=Text(content=value[idx : idx + chunk_size])) - for idx in range(0, len(value), chunk_size) - ] - - -class Mention(BaseModel): - pass # TODO - - class SelectOption(BaseModel): id: Optional[str] name: str @@ -296,3 +249,73 @@ def value(self): elif self.type == "date": if self.date is not None: return self.date.value + + +class AnnotationObject(BaseModel): + bold: bool + italic: bool + strikethrough: bool + underline: bool + code: bool + color: NotionExtendedColorEnum + + +class TextLinkObject(BaseModel): + type: Optional[str] = "url" + url: str + + +class TextObject(BaseModel): + content: str + link: Optional[TextLinkObject] + + +class PageReferenceObject(BaseModel): + id: str + + +class LinkPreviewMentionObject(BaseModel): + url: str + + +class MentionObject(BaseModel): + type: str + user: Optional[UserObject] + page: Optional[PageReferenceObject] + database: Optional[PageReferenceObject] + date: Optional[DateObject] + link_preview: Optional[LinkPreviewMentionObject] + + +class EquationObject(BaseModel): + expression: str + + +class BaseRichTextObject(BaseModel): + plain_text: Optional[str] + # TODO: The Optional[plain_text] is used when creating property values + href: Optional[str] = None + annotations: Optional[AnnotationObject] = None + type: Optional[RichTextTypeEnum] + + @property + def value(self): + return self.plain_text + + +class RichTextObject(BaseRichTextObject): + text: Optional[TextObject] + mention: Optional[MentionObject] + equation: Optional[EquationObject] + + @classmethod + def from_value(cls, value: str): + return cls(text=TextObject(content=value)) + + @classmethod + def encode_string(cls, value: str) -> List["RichTextObject"]: + chunk_size = RICH_TEXT_CONTENT_MAX_LENGTH + return [ + cls(text=TextObject(content=value[idx : idx + chunk_size])) + for idx in range(0, len(value), chunk_size) + ] \ No newline at end of file diff --git a/src/notion_df/values.py b/src/notion_df/values.py index 041c0f1..d53ce4c 100644 --- a/src/notion_df/values.py +++ b/src/notion_df/values.py @@ -10,7 +10,7 @@ from pandas.api.types import is_array_like from notion_df.base import ( - RichText, + RichTextObject, SelectOption, DateObject, RelationObject, @@ -41,36 +41,36 @@ def query_dict(self): class TitleValues(BasePropertyValues): - title: List[RichText] + title: List[RichTextObject] @property def value(self) -> Optional[str]: return ( None if len(self.title) == 0 - else " ".join([text.plain_text for text in self.title]) + else " ".join([text.value for text in self.title]) ) @classmethod def from_value(cls, value): - return cls(title=RichText.encode_string(value)) + return cls(title=RichTextObject.encode_string(value)) # TODO: Rethink whether we should split input string to multiple elements in the list class RichTextValues(BasePropertyValues): - rich_text: List[RichText] + rich_text: List[RichTextObject] @property def value(self) -> Optional[str]: return ( None if len(self.rich_text) == 0 - else " ".join([text.plain_text for text in self.rich_text]) + else " ".join([text.value for text in self.rich_text]) ) @classmethod def from_value(cls, value: str): - return cls(rich_text=RichText.encode_string(value)) + return cls(rich_text=RichTextObject.encode_string(value)) class NumberValues(BasePropertyValues): From c7fcbeae8b3f69960cd0c4c87dc371aaaa708d6b Mon Sep 17 00:00:00 2001 From: Shannon Shen <22512825+lolipopshock@users.noreply.github.com> Date: Sun, 13 Feb 2022 20:14:08 -0500 Subject: [PATCH 2/2] Add tests --- tests/test_base.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/test_base.py b/tests/test_base.py index 4f90f07..9232be3 100644 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -143,4 +143,13 @@ def test_long_string(): upload(df[:1], NOTION_LONG_STRING_DF, api_key=NOTION_API_KEY) df_new = download(NOTION_LONG_STRING_DF, api_key=NOTION_API_KEY) - assert len(df.iloc[0,1]) == 7721 \ No newline at end of file + # assert len(df_new.iloc[0,1]) == 7721 + # This might not be true -- understand why? + +def test_rich_text(): + NOTION_RICH_TEXT_DF = os.environ.get("NOTION_RICH_TEXT_DF") + + if not NOTION_RICH_TEXT_DF or not NOTION_API_KEY: + pytest.skip("API key not provided") + + df = download(NOTION_RICH_TEXT_DF, api_key=NOTION_API_KEY) \ No newline at end of file