Local-Booru/system_module/tasks.py

import os
import json
from pathlib import Path

def find_file(filename, file_folders):
    for folder in file_folders:
        folder_path = Path(folder)
        if not folder_path.exists():
            continue
        for file_path in folder_path.rglob(filename):
            return folder_path
    return None

def get_size_format(b, factor=1024, suffix="B"):
    for unit in ["", "K", "M", "G", "T"]:
        if b < factor:
            return f"{b:.2f} {unit}{suffix}"
        b /= factor
    return f"{b:.2f} P{suffix}"

def get_dir_size(path: str):
    total = 0
    if os.path.exists(path):
        for f in os.listdir(path):
            fp = os.path.join(path, f)
            if os.path.isfile(fp):
                total += os.path.getsize(fp)
    return total

def get_file_list_size(file_list: list[str]):
    total = 0
    for fp in file_list:
        if os.path.isfile(fp):
            total += os.path.getsize(fp)
    return total

def get_size_edge(type_searсh: bool, file_list: list[str]):
    if type_searсh:
        return min(file_list, key=os.path.getsize)
    else:
        return max(file_list, key=os.path.getsize)

def calculate_analytic(database, config):
    all_db = database.get_all()
    explicit_db = database.get_search(["e"], [])
    sensitive_db = database.get_search(["s"], [])
    questionable_db = database.get_search(["q"], [])
    general_db = database.get_search(["g"], [])
    rating_list_file_path = {"explicit":[], "sensitive":[], "questionable":[], "general":[]}
    type_list_file_path = {"images":[], "videos":[], "gifs":[]}
    files_not_found = 0
    for one_post in all_db:
        file_path = find_file(one_post["local_filename"], config["file"]["repos"])
        if file_path != None:
            rating_list_file_path[one_post["rating"]].append(file_path)
            if one_post["is_gif"] == 1:
                type_list_file_path["gifs"].append(file_path)
            elif one_post["is_video"] == 1:
                type_list_file_path["videos"].append(file_path)
            else:
                type_list_file_path["images"].append(file_path)
        else:
            files_not_found += 1
    file_size_dict = {"explicit":get_file_list_size(rating_list_file_path["explicit"]), 
                        "sensitive":get_file_list_size(rating_list_file_path["sensitive"]), 
                        "questionable":get_file_list_size(rating_list_file_path["questionable"]), 
                        "general":get_file_list_size(rating_list_file_path["general"])
                    }
    type_size_dict = {"images":get_file_list_size(type_list_file_path["images"]),
                        "videos":get_file_list_size(type_list_file_path["videos"]), 
                        "gifs":get_file_list_size(type_list_file_path["gifs"])}
    ratings_stats = {
        "explicit": {
            "count": len(explicit_db),
            "size_formatted": get_size_format(file_size_dict["explicit"])
        },
        "sensitive": {
            "count": len(sensitive_db),
            "size_formatted": get_size_format(file_size_dict["sensitive"])
        },
        "questionable": {
            "count": len(questionable_db),
            "size_formatted": get_size_format(file_size_dict["questionable"])
        },
        "general": {
            "count": len(general_db),
            "size_formatted": get_size_format(file_size_dict["general"])
        }
    }
    total_file_data_size = file_size_dict["general"]+file_size_dict["questionable"]+file_size_dict["sensitive"]+file_size_dict["explicit"]
    total_thumb_size = get_dir_size(config[["file"]["thumb"]])
    total_log_size = get_file_list_size(["logs.log"])
    total_db_size = get_file_list_size([os.path.join("data", "db.sqlite")])
    weight_stats_smallest = get_size_edge(True, type_list_file_path["gifs"]+type_list_file_path["images"]+type_list_file_path["videos"])
    weight_stats_largest = get_size_edge(False, type_list_file_path["gifs"]+type_list_file_path["images"]+type_list_file_path["videos"])
    result = {
        "total_images_analyzed": len(all_db),
        "files_not_found":files_not_found,
        "average_weight": {
            "bytes": round(total_file_data_size/(len(rating_list_file_path["explicit"])+len(rating_list_file_path["sensitive"])+len(rating_list_file_path["questionable"])+len(rating_list_file_path["general"])), 2),
            "formatted": get_size_format(round(total_file_data_size/(len(rating_list_file_path["explicit"])+len(rating_list_file_path["sensitive"])+len(rating_list_file_path["questionable"])+len(rating_list_file_path["general"])), 2))
        },
        "storage_stats": {
            "folders": {
                "data": get_size_format(total_file_data_size),
                "db": get_size_format(total_db_size),
                "thumb": get_size_format(total_thumb_size),
                "logs": get_size_format(total_log_size)
            },
            "types": {
                "images": get_size_format(type_size_dict["images"]),
                "videos": get_size_format(type_size_dict["videos"]),
                "gifs": get_size_format(type_size_dict["gifs"])
            },
            "grand_total": get_size_format(total_file_data_size+total_thumb_size+total_log_size+total_db_size)
        },
        "ratings_stats": ratings_stats,
        "resolution_stats": {
            "largest": {
                "filename": max_res["filename"],
                "md5": max_res["md5"],
                "resolution": f'{max_res["width"]}x{max_res["height"]}'
            },
            "smallest": {
                "filename": min_res["filename"],
                "md5": min_res["md5"],
                "resolution": f'{min_res["width"]}x{min_res["height"]}'
            }
        },
        "weight_stats": {
            "largest": {
                "filename": weight_stats_largest,
                "md5": max_weight["md5"],
                "size_formatted": get_size_format(max_weight["size"])
            },
            "smallest": {
                "filename": weight_stats_smallest,
                "md5": min_weight["md5"],
                "size_formatted": get_size_format(min_weight["size"])
            }
        },
        "top_tag":database.get_top_tag("")
    }
    with open(os.path.join(config["file"]["data"], "analytic.json"), 'w', encoding='utf-8') as out_f:
        json.dump(result, out_f, indent=4, ensure_ascii=False)