0.0.6s

2026-06-02 00:15:38 +03:00
parent 2a103d246b
commit 627b374b7f
8 changed files with 436 additions and 84 deletions
@@ -0,0 +1,104 @@
+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)