scrubber/2_filter.py
2024-11-20 06:52:59 -08:00

69 lines
1.8 KiB
Python

import time
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, List
import psycopg
from com import (FilterableNote, Visibility, eval_config, parse_graph,
progressbar, FilterAction)
config = eval_config()
conn: psycopg.Connection = config["connect"]()
criteria: Callable[[FilterableNote], FilterAction] = config["criteria"]
intermediate = parse_graph()
def transform(entry: dict) -> FilterableNote:
time.sleep(0.0001)
note = conn.execute(
'select "createdAt", reactions, "renoteCount", visibility, cw from note where id = %s',
[entry["id"]],
).fetchone()
if note is None:
return None # part of thread disappeared during processing
when, reactions, renotes, visibility, cw = note
replies = [transform(intermediate[reply]) for reply in entry["replies"]]
quotes = [transform(intermediate[quote]) for quote in entry["quotes"]]
if None in replies or None in quotes:
return None # bubble up, buttercup
return FilterableNote(
entry["id"],
"self" in entry["flags"],
replies,
quotes,
when.astimezone(),
sum(reactions.values()),
renotes,
Visibility.from_db(visibility),
cw,
)
root_count = 0
for entry in intermediate.values():
if "root" in entry["flags"]:
root_count += 1
pb = progressbar.ProgressBar(
0,
root_count,
prefix="processing ",
)
targets = []
for entry in intermediate.values():
if "root" not in entry["flags"]:
continue
transformed = transform(entry)
if transformed is None:
continue # we'll get to it next cycle
action = criteria(transformed)
if action != FilterAction.Ignore:
targets.append(f"{entry['id']} {action.value}")
pb.increment()
pb.finish()
Path("filtered.list").write_text("\n".join(targets))