Hell13Cat
92 lines
2.8 KiB
Python
92 lines
2.8 KiB
Python
from tinydb import TinyDB, Query
|
|
import os
|
|
import json
|
|
from collections import Counter
|
|
import gelbooru
|
|
from datetime import datetime
|
|
|
|
class DB:
|
|
def __init__(self):
|
|
db_path = "datas"
|
|
if not os.path.exists(db_path):
|
|
os.makedirs(os.path.dirname(db_path), exist_ok=True)
|
|
json.dump({}, open(db_path+"/db.json", "w", encoding="utf-8"))
|
|
self.db = TinyDB(db_path+"/db.json")
|
|
self.File = Query()
|
|
|
|
def add(self, data_add):
|
|
self.db.insert(data_add)
|
|
|
|
def add_raw(self, data_add):
|
|
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)
|
|
pass
|
|
|
|
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, ratings, tags):
|
|
ratings
|
|
results = self.db.search(
|
|
self.File.rating.one_of(self.ratings) &
|
|
self.File.tags.any(self.tags)
|
|
)
|
|
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):
|
|
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(prefix)
|
|
]
|
|
results.sort(key=lambda x: x["count"], reverse=True)
|
|
top_10 = results[:10]
|
|
return top_10
|
|
|
|
def edit_data(self, id_file, edit_data):
|
|
self.db.update(self.edit_data, self.File.md5 == self.id_file) |