Source code for mcot.org.database

from pathlib import Path
from typing import List, Set, Optional, Dict, Tuple
from dataclasses import dataclass, field
import pandas as pd
import sqlite3
import os.path as op
import re

@dataclass(unsafe_hash=True)
class RoamFile:
    filename: str
    hash: str
    contains: List["RoamNode"] = field(default_factory=list, hash=False)

    @property
    def basename(self, ) -> str:
        return op.basename(self.filename)

    def url(self, extension='md'):
        """How will target be availabe

        By default will return the filename with .org replaced by the extension, except for:

            - any date prefix is removed
            - if output is HTML introduction_to_my_brain will be replaced by index.html
            - identifier is added to the end

        :param extension: filename extension to use
        :return: url
        """
        if self.roam_key is not None and self.roam_key.startswith('cite:'):
            return self.roam_key[5:] + '.' + extension
        return self.title.replace(':', ';').replace('/', ', ').replace('-', 'in') + '.' + extension

    @property
    def tags(self, ):
        return {t for c in self.contains for t in c.tags}

    @property
    def top_level(self, ) -> "RoamNode":
        return [c for c in self.contains if c.level == 0][0]

    @property
    def roam_key(self, ):
        return self.top_level.roam_key

    def __getattr__(self, key: str) -> bool:
        if key in ('publish', 'book', 'person', 'private', 'web'):
            return key in self.tags
        raise AttributeError

    @property
    def links_from(self, ):
        return {n for node in self.contains for n in node.links_from}

    @property
    def links_to(self, ):
        return {n for node in self.contains for n in node.links_to}

    @property
    def external_refs(self, ):
        return [n for node in self.contains for n in node.external_refs]

    @property
    def title(self, ):
        return self.top_level.title

@dataclass
class ExternalLink:
    dest: str
    type: str
    source: "RoamNode"

@dataclass(unsafe_hash=True)
class RoamNode:
    title: str
    file: RoamFile
    level: int
    pos: int
    identifier: str
    properties: Dict[str, str]
    tags: List[str] = field(default_factory=list, hash=False)
    links_to: Set["RoamNode"] = field(default_factory=set, repr=False, hash=False)
    links_from: Set["RoamNode"] = field(default_factory=set, repr=False, hash=False)
    external_refs: List[ExternalLink] = field(default_factory=list, repr=False, hash=False)

    def url(self, extension='md'):
        file_url = self.file.url(extension)
        if self.level == 0:
            return file_url
        else:
            return file_url + '#' + self.title

    @property
    def roam_key(self, ) -> Optional[str]:
        """Returns the key stored in #+ROAM_REFS (None if not avaialable)"""
        return self.properties.get('ROAM_REFS', None)

    @property
    def ref(self, ) -> bool:
        return self.roam_key is not None and self.roam_key.startswith('cite:')

    def __getattr__(self, key: str) -> bool:
        if key in ('publish', 'book', 'person', 'private', 'web'):
            return key in self.tags
        raise AttributeError


def default_database():
    database = Path.home() / ".emacs.d" / ".local" / "etc" / "org-roam.db"
    if not database.is_file():
        raise ValueError("Default org-roam database not found at {database}. Please provide database path.")
    return database


def load(database) -> Tuple[List[RoamFile], List[RoamNode]]:
    """
    Loads in all Files from database

    :param database: org-roam database
    :return: list of all org-roam files in database
    """
    con = sqlite3.connect(database)

    tables = {
        name: pd.read_sql_query(f"SELECT * from {name}", con)
        for name in ("tags", "files", "refs", "links", "nodes")
    }
    con.close()

    files_list: List[RoamFile] = []
    files_dict: Dict[str, RoamFile] = {}
    for _, row in tables['files'].iterrows():
        new_file = RoamFile(
            filename=row['file'][1:-1], 
            hash=row['hash'][1:-1], 
        )
        files_list.append(new_file)
        files_dict[new_file.filename] = new_file


    nodes_list: List[RoamNode] = []
    nodes_dict: Dict[str, RoamNode] = {}
    for _, row in tables['nodes'].iterrows():
        new_node = RoamNode(
            title=row['title'][1:-1],
            level=row['level'],
            pos=row['pos'],
            file=files_dict[row['file'][1:-1]],
            identifier=row['id'][1:-1],
            properties=read_properties(row['properties']),
        )
        new_node.file.contains.append(new_node)
        nodes_list.append(new_node)
        nodes_dict[new_node.identifier] = new_node

    for _, row in tables['tags'].iterrows():
        nodes_dict[row['node_id'][1:-1]].tags.append(row['tag'][1:-1])

    for _, row in tables['links'].iterrows():
        source = nodes_dict[row['source'][1:-1]]
        if row['type'] == 'id':
            dest = nodes_dict[row['dest'][1:-1]]
            source.links_to.add(dest)
            dest.links_from.add(source)
        else:
            source.external_refs.append(ExternalLink(type=row['type'][1:-1], dest=row['dest'][1:-1], source=source))
    return files_list, nodes_list


def read_properties(text):
    return {key: value for key, value in re.findall('\("(.*?)" \. "(.*?)"\)', text)}