Toggle Light / Dark / Auto color theme

Toggle table of contents sidebar

Source code for remarking.models

""" Contains SqlAlchemy models and their helper methods """
import datetime
import hashlib
import typing as T
from typing import Dict

from dateutil import parser as date_parser
from rmapy import document as rmapy_document
from sqlalchemy import (TIMESTAMP, Boolean, Column, ForeignKey, Integer,
                        String, Text)
from sqlalchemy.orm import declarative_base

Base = declarative_base()
C = T.TypeVar('C', bound='ModelMixIn')  # pylint: disable=C0103


class ModelMixIn():
    """ MixIn for common functionality used for SqlAlchemy models. """

    @classmethod
    def from_dict(cls: T.Type[C], data: Dict[str, T.Any]) -> C:
        """ Construct model from a dictionary.

            data should have keys that match model names exactly.
        """
        columns = [column.name for column in cls.__table__.columns]  # type: ignore
        constructor_args = {}
        for key, value in data.items():
            if key in columns:
                constructor_args[key] = value
        return cls(**constructor_args)  # type: ignore

    def to_dict(self) -> Dict[str, T.Any]:
        """ Return a dictionary that represents the model. """
        return {c.name: getattr(self, c.name) for c in self.__table__.columns}  # type: ignore

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}({str(self.to_dict())})"


class Document(Base, ModelMixIn):
    """ Represents the current state of a document originating from reMarkable cloud """

    __tablename__ = "document"
    id: str = Column(String(256), nullable=False, primary_key=True)
    """ Primary key for document. This is a UUID generated by the reMarkable tablet. """
    version: int = Column(Integer, nullable=False)
    """ The version of the document according to the reMarkable cloud. """
    modified_client: datetime.datetime = Column(TIMESTAMP, nullable=False)
    """ The unix timestamp for the last time this document was modified. """
    type: str = Column(String(256), nullable=False)
    """ The type of the document. """
    name: str = Column(Text, nullable=False)
    """ The name of the document as visible on the reMarkable tablet. """
    current_page: int = Column(Integer, nullable=False)
    """ The current page the document is opened on. """
    bookmarked: bool = Column(Boolean, nullable=False)
    """ Indicate if the document is bookmarked. """
    parent: str = Column(String(256), nullable=False)
    """ The parent of the document. This is usually a folder ID. """

    def equal(self, other: 'Document') -> bool:
        """ Check for equality with other documents. """
        return self.id == other.id

    @classmethod
    def from_cloud_document(cls, cloud_document: rmapy_document.Document) -> 'Document':
        """ Create a Document model from a reMarkable cloud document.  """
        cloud_document_metadata = cloud_document.to_dict()
        cloud_document_mapping = {
            "ID": "id",
            "Version": "version",
            "Message": "message",
            "Success": "success",
            "ModifiedClient": "modified_client",
            "Type": "type",
            "VissibleName": "name",
            "CurrentPage": "current_page",
            "Bookmarked": "bookmarked",
            "Parent": "parent"

        }
        cleaned_metadata = {}
        for key, value in cloud_document_metadata.items():
            if key not in cloud_document_mapping:
                continue
            new_key = cloud_document_mapping[key]
            if key == "ModifiedClient":
                cleaned_metadata[new_key] = date_parser.parse(value).replace(tzinfo=None)
            else:
                cleaned_metadata[new_key] = value

        return cls.from_dict(cleaned_metadata)

    def to_metadata_dict(self) -> Dict[str, T.Any]:
        """ Return a dictionary that matches the .content file used by the reMarkable. """
        return {
            "deleted": False,
            "lastModified": self.modified_client.timestamp() if self.modified_client else None,
            "lastOpenedPage": self.current_page,
            "metadatamodified": False,
            "modified": False,
            "parent": self.parent,
            "pinned": False,
            "synced": True,
            "type": self.type,
            "version": self.version,
            "visibleName": self.name
        }


class Highlight(Base, ModelMixIn):
    """ Represents a single highlight for a Document. """
    __tablename__ = "highlight"
    hash: str = Column(String(256), primary_key=True, nullable=False)
    """ Primary for a highlight. This is a hash of the document_id and text. """
    document_id: str = Column(String(256), ForeignKey('document.id'), nullable=False)
    """ The document that the highlight is associated with """
    text: str = Column(Text)
    """ The text of the highlight """
    page_number: int = Column(Integer)
    """ The page number on which the highlight is located. """
    extracted_at: datetime.datetime = Column(TIMESTAMP)  # type: ignore
    """ A unix timestamp of when the highlight was extracted. """
    extraction_method: str = Column(String(256), nullable=True)
    """ What method was used to perform the highlight extraction. """

    @classmethod
    def create_highlight(cls, doc_id: str, text: str, page_number: int, extraction_method: str) -> 'Highlight':
        """ Create a Highlight.
        This should be used in place of the default constructor as it properly constructs the highlight hash.

        :param doc_id: The id of the document this highlight is associated with.
        :param text: The text of the highlight.
        :param page_number: The page number where the highlight is located.
        :param extract_method: The extraction method.

        :return: An instance of highlight for these values.
        """
        return Highlight(
            document_id=doc_id,
            text=text,
            hash=hashlib.sha224((text + doc_id).encode()).hexdigest(),
            page_number=page_number,
            extraction_method=extraction_method,
            extracted_at=datetime.datetime.now()
        )

    def equal(self, other: 'Highlight') -> bool:
        """ Check for equality with other Highlights. """
        return (self.hash == other.hash and
                self.text == other.text and
                self.document_id == other.document_id and
                self.page_number == other.page_number)