2026-06-06 12:34:50 +00:00
1 changed files with 0 additions and 104 deletions
@@ -1,104 +0,0 @@
 from tinydb import TinyDB, Query
 import os
 import json
 from collections import Counter
 from system_module import gelbooru
 from datetime import datetime
 class DB:
    def __init__(self, debug: bool):
        if debug: 
            db_path = "data_test" 
        else: 
            db_path = "data"
        if not os.path.exists(os.path.join(db_path, "db.json")):
            json.dump({}, open(os.path.join(db_path, "db.json"), "w", encoding="utf-8"))
        self.db = TinyDB(db_path+"/db.json")
        self.File = Query()
    def add(self, data_add: dict):
        self.db.insert(data_add)
    def add_raw(self, data_add: dict):
        is_video = False
        is_gif = False
        if data_add["image"].split(".")[-1] in ["webm", "mp4", "avi"]:
            is_video = True
        if data_add["image"].split(".")[-1] == "gif":
            is_gif = True
        add_data = {
            "md5": data_add["md5"],
            "local_filename": data_add["image"],
            "is_video": is_video,
            "is_gif": is_gif,
            "tags": data_add["tags"].split(" "),
            "rating": data_add["rating"],
            "score": data_add["score"],
            "timestamp": data_add["created_at"],
            "display_date": gelbooru.gelbooru_date_parse(data_add["created_at"]),
            "source": data_add["source"],
            "width": data_add["width"],
            "height": data_add["height"]
        }
        self.db.insert(add_data)
    def get_all(self):
        sorted_records = sorted(
            self.db.all(),
            key=lambda x: datetime.strptime(
                x["timestamp"],
                "%a %b %d %H:%M:%S %z %Y"
            ),
            reverse=True
        )
        return sorted_records
    def get_search(self, rating_tmp: list[str], tags: list[str]):
        mapping_rating = {
            "e": "explicit",
            "g": "general",
            "s": "sensitive",
            "q": "questionable"
        }
        rating = [mapping_rating.get(x, x) for x in rating_tmp]
        if rating:
            query = self.File.rating.one_of(rating)
        if tags != [""]:
            tag_query = self.File.tags.test(lambda t: set(tags).issubset(set(t or [])))
            query = tag_query if query is None else query & tag_query
        results = self.db.search(query) if query else self.db.all()
        sorted_records = sorted(
            results,
            key=lambda x: datetime.strptime(
                x["timestamp"],
                "%a %b %d %H:%M:%S %z %Y"
            ),
            reverse=True
        )
        return sorted_records
    def get_unknow(self):
        results = self.db.search(self.File.tags == [])
        return results
    def get_autocomplete(self, tag_slise: str):
        counter = Counter()
        for record in self.db.all():
            counter.update(record.get("tags", []))
        results = [
            {
                "tag": tag,
                "count": count
            }
            for tag, count in counter.items()
            if tag.lower().startswith(tag_slise)
        ]
        results.sort(key=lambda x: x["count"], reverse=True)
        top_10 = results[:10]
        return top_10
    def edit_data(self, id_file: str, edit_data: dict):
        self.db.update(edit_data, self.File.md5 == id_file)
    def get_from_id(self, id_file: str):
        return self.db.get(self.File.md5 == id_file)