scrubber/com.py

import sys
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Callable, Dict, List

try:
    import progressbar2 as progressbar
except ImportError:
    import progressbar


class Visibility(Enum):
    public = 1
    unlisted = 2
    followers = 3
    direct = 4

    @classmethod
    def from_db(cls, raw: str) -> "Visibility":
        match raw:
            case "public": return cls.public
            case "home": return cls.unlisted
            case "followers": return cls.followers
            case "specified": return cls.direct
            case _: raise ValueError(f"unknown visibility `{raw}`")
    
    def code(self) -> str:
        match self:
            case self.public: return "p"
            case self.unlisted: return "u"
            case self.followers: return "f"
            case self.direct: return "d"


@dataclass
class FilterableNote:
    id: str
    mine: bool
    replies: List["FilterableNote"]
    quotes: List["FilterableNote"]
    when: datetime
    reactions: int
    renotes: int
    visibility: Visibility

    def thread(self) -> List["FilterableNote"]:
        acc = []
        for reply in self.replies:
            acc += reply.thread()
        for quote in self.quotes:
            acc += quote.thread()
        acc.append(self)
        return acc

    def thread_self(self) -> List["FilterableNote"]:
        acc = []
        for reply in self.replies:
            acc += reply.thread_self()
        for quote in self.quotes:
            acc += quote.thread_self()
        if self.mine:
            acc.append(self)
        return acc

    def to_dict(self):
        return {
            "id": self.id,
            "mine": self.mine,
            "replies": [note.to_dict() for note in self.replies],
            "quotes": [note.to_dict() for note in self.quotes],
            "when": self.when.isoformat(),
            "reactions": self.reactions,
            "renotes": self.renotes,
        }


def eval_config() -> dict:
    print("configuring")
    config = {}
    exec(Path(sys.argv[1]).read_text(), config)
    return config


def parse_graph() -> Dict[str, dict]:
    print("parsing graph")
    graph = {}
    for line in Path("graph.db").read_text().splitlines():
        id, replies, quotes, flags = line.split("\t")
        graph[id] = {
            "id": id,
            "replies": replies.split(",") if len(replies) > 0 else [],
            "quotes": quotes.split(",") if len(quotes) > 0 else [],
            "flags": flags.split(",") if len(flags) > 0 else [],
        }
    return graph