home-automation-backend/src/cloud_util/notion.py

from __future__ import annotations

from dataclasses import asdict, dataclass, field

from notion_client import AsyncClient as Client


@dataclass
class Text:
    content: str


@dataclass
class RichText:
    type: str
    href: str | None = None


@dataclass
class RichTextText(RichText):
    type: str = "text"
    text: Text = field(default_factory=lambda: Text(content=""))


class NotionAsync:
    def __init__(self, token: str) -> None:
        self._client = Client(auth=token)

    def update_token(self, token: str) -> None:
        self._client.aclose()
        self._client = Client(auth=token)

    async def get_block(self, block_id: str) -> dict:
        return await self._client.blocks.retrieve(block_id=block_id)

    async def get_block_children(self, block_id: str, start_cursor: str | None = None, page_size: int = 100) -> dict:
        return await self._client.blocks.children.list(
            block_id=block_id,
            start_cursor=start_cursor,
            page_size=page_size,
        )

    async def block_is_table(self, block_id: str) -> bool:
        block: dict = await self.get_block(block_id=block_id)
        return block["type"] == "table"

    async def get_table_width(self, table_id: str) -> int:
        table = await self._client.blocks.retrieve(block_id=table_id)
        return table["table"]["table_width"]

    async def append_table_row_text(self, table_id: str, text_list: list[str], after: str | None = None) -> None:
        cells: list[RichText] = []
        for content in text_list:
            cells.append([asdict(RichTextText(text=Text(content)))])  # noqa: PERF401
        await self.append_table_row(table_id=table_id, cells=cells, after=after)

    async def append_table_row(self, table_id: str, cells: list[RichText], after: str | None = None) -> None:
        if not await self.block_is_table(table_id):
            return
        table_width = await self.get_table_width(table_id=table_id)
        if table_width != len(cells):
            return
        children = [
            {
                "object": "block",
                "type": "table_row",
                "table_row": {
                    "cells": cells,
                },
            },
        ]
        if after is None:
            await self._client.blocks.children.append(
                block_id=table_id,
                children=children,
            )
        else:
            await self._client.blocks.children.append(
                block_id=table_id,
                children=children,
                after=after,
            )