Upload df with children (#32)

* better organization of blocks.py

* Allow upload with children

Author: lolipopshock

src/notion_df/agent.py

@@ -11,8 +11,8 @@
 
 from notion_df.values import PageProperties, PageProperty
 from notion_df.configs import DatabaseSchema, NON_EDITABLE_TYPES
-from notion_df.utils import is_uuid
-from notion_df.blocks import parse_blocks
+from notion_df.utils import is_uuid, flatten_dict
+from notion_df.blocks import parse_blocks, BaseNotionBlock
 
 API_KEY = None
 NOT_REVERSE_DATAFRAME = -1
@@ -235,20 +235,36 @@ def create_database(
     return response
 
 
-def upload_row_to_database(row, database_id, schema, client) -> Dict:
+def upload_row_to_database(row, database_id, schema, children, client) -> Dict:
 
     properties = PageProperty.from_series(row, schema).query_dict()
-    response = client.pages.create(
-        parent={"database_id": database_id}, properties=properties
-    )
+    if children:
+        if not isinstance(children, list):
+            children = [children]
+        for cid in range(len(children)):
+            if isinstance(children[cid], BaseNotionBlock):
+                children[cid] = flatten_dict(children[cid].dict())
+                
+        response = client.pages.create(
+            parent={"database_id": database_id}, properties=properties, children=children
+        )
+    else:
+        response = client.pages.create(
+            parent={"database_id": database_id}, properties=properties,
+        )
     return response
 
 
-def upload_to_database(df, databse_id, schema, client, errors) -> List[Dict]:
+def upload_to_database(df, databse_id, schema, client, errors, children) -> List[Dict]:
     all_response = []
-    for _, row in df[::NOT_REVERSE_DATAFRAME].iterrows():
+    if children is not None:
+        assert len(children) == len(df)
+        children = children[::NOT_REVERSE_DATAFRAME]
+
+    for idx, (_, row) in enumerate(df[::NOT_REVERSE_DATAFRAME].iterrows(), ):
         try:
-            response = upload_row_to_database(row, databse_id, schema, client)
+            child = children[idx] if children is not None else None
+            response = upload_row_to_database(row, databse_id, schema, child, client)
             all_response.append(response)
         except Exception as e:
             if errors == "strict":
@@ -277,6 +293,7 @@ def upload(
     errors: str = "strict",
     resolve_relation_values: bool = False,
     create_new_rows_in_relation_target: bool = False,
+    children: List[Union[Dict, BaseNotionBlock]] = None,
     return_response: bool = False,
     *,
     api_key: str = None,
@@ -317,6 +334,9 @@ def upload(
                     subsequent rows.
                 3. "warn": print the error message and continue uploading
             Defaults to "strict".
+        children (List[Union[Dict, BaseNotionBlock]], optional):
+            The corresponding children of the uploaded Notion page. It should be
+            a list of the same length as the dataframe.
         resolve_relation_values (bool, optional):
             If `True`, notion-df assumes the items in any relation columns
             are not notion object ids, but the value of the corresponding 
@@ -432,7 +452,7 @@ def upload(
                         lambda row: [obj_string_to_id[ele] for ele in row if ele in obj_string_to_id]
                     )
 
-    response = upload_to_database(df, databse_id, schema, client, errors)
+    response = upload_to_database(df, databse_id, schema, client, errors, children)
 
     print(f"Your dataframe has been uploaded to the Notion page: {notion_url} .")
     if return_response:

src/notion_df/blocks.py

@@ -26,25 +26,9 @@ class ParentObject(BaseModel):
     block_id: Optional[str]
 
 
-class BaseNotionBlock(BaseModel):
-    object: str = "block"
-    parent: ParentObject
-    id: Optional[str]
-    type: Optional[str]
-    created_time: str
-    # created_by
-    last_edited_time: str
-    # created_by
-    has_children: bool
-    archived: bool
-    type: str
-
-    @property
-    def children(self):
-        return self.__getattribute__(self.type).children
-
-    def set_children(self, value: Any):
-        self.__getattribute__(self.type).children = value
+# BaseClasses
+class BaseAttributes(BaseModel):
+    pass
 
 
 class BaseAttributeWithChildren(BaseModel):
@@ -75,60 +59,89 @@ class ToDoBlockAttributes(BaseAttributeWithChildren):
     checked: Optional[bool]
 
 
-class CodeBlockAttributes(BaseModel):
+class CodeBlockAttributes(BaseAttributes):
     rich_text: List[RichTextObject]
     caption: Optional[List[RichTextObject]]
     language: Optional[str]  # TODO: it's actually an enum
 
 
-class ChildPageAttributes(BaseModel):
+class ChildPageAttributes(BaseAttributes):
     title: List[RichTextObject]
 
 
-class EmbedBlockAttributes(BaseModel):
+class EmbedBlockAttributes(BaseAttributes):
     url: str
 
 
-class ImageBlockAttributes(FileObject):
+class ImageBlockAttributes(BaseAttributes, FileObject):
     caption: Optional[List[RichTextObject]]
     # This is not listed in the docs, but it is in the API response (Nov 2022)
 
-class VideoBlockAttributes(BaseModel):
+
+class VideoBlockAttributes(BaseAttributes):
     video: FileObject
 
 
-class FileBlockAttributes(BaseModel):
+class FileBlockAttributes(BaseAttributes):
     file: FileObject
     caption: Optional[List[RichTextObject]]
 
 
-class PdfBlockAttributes(BaseModel):
+class PdfBlockAttributes(BaseAttributes):
     pdf: FileObject
 
 
-class BookmarkBlockAttributes(BaseModel):
+class BookmarkBlockAttributes(BaseAttributes):
     url: str
     caption: Optional[List[RichTextObject]]
 
 
-class EquationBlockAttributes(BaseModel):
+class EquationBlockAttributes(BaseAttributes):
     expression: str
 
 
-class TableOfContentsAttributes(BaseModel):
+class TableOfContentsAttributes(BaseAttributes):
     color: Optional[NotionExtendedColorEnum]
 
 
-class LinkPreviewAttributes(BaseModel):
+class LinkPreviewAttributes(BaseAttributes):
     url: str
 
 
-class LinkToPageAttributes(BaseModel):
+class LinkToPageAttributes(BaseAttributes):
     type: str
     page_id: Optional[str]
     database_id: Optional[str]
 
 
+ATTRIBUTES_MAPPING = {
+    _cls.__name__: _cls
+    for _cls in BaseAttributes.__subclasses__()
+    + BaseAttributeWithChildren.__subclasses__()
+}
+
+
+class BaseNotionBlock(BaseModel):
+    object: str = "block"
+    parent: Optional[ParentObject]
+    id: Optional[str]
+    type: Optional[str]
+    created_time: Optional[str]
+    # created_by
+    last_edited_time: Optional[str]
+    # created_by
+    has_children: Optional[bool]
+    archived: Optional[bool]
+    type: str
+
+    @property
+    def children(self):
+        return self.__getattribute__(self.type).children
+
+    def set_children(self, value: Any):
+        self.__getattribute__(self.type).children = value
+
+
 class ParagraphBlock(BaseNotionBlock):
     type: str = "paragraph"
     paragraph: TextBlockAttributes