utils/stats.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# ファイルの統計情報を取得する
# Get statistics of files

import os
import sys
from collections import Counter
import gettext
import locale
from pathlib import Path
import logging

# Set up i18n
def setup_i18n():
    """Set up internationalization"""
    try:
        # Get the current locale
        locale.setlocale(locale.LC_ALL, '')
        current_locale = locale.getlocale()[0]
        locale_path = Path(__file__).parent / 'locales'
        
        # Try to load translations for current locale
        trans = gettext.translation('stats', locale_path, languages=[current_locale])
        trans.install()
        return trans.gettext
    except:
        # Fallback to default English if translation fails
        return gettext.gettext

# Initialize translation
_ = setup_i18n()

# ANSI color codes
RED = "\033[91m"
GREEN = "\033[92m"
ORANGE = "\033[93m"
BLUE = "\033[94m"
MAGENTA = "\033[95m"
CYAN = "\033[96m"
RESET = "\033[0m"

EXT2COLOR = {
    "jxl": CYAN,
    "png": MAGENTA,
    "jpg": RED,
    "jpeg": RED,
    "webp": MAGENTA,
    "caption": BLUE,
    "e621": BLUE,
    "txt": BLUE,
    "wd": BLUE,
    "json": ORANGE,
    "md": ORANGE,
    "tags": BLUE,
}
EXT2ORDER = {ext: i for i, ext in enumerate(EXT2COLOR.keys())}
SORT_KEYS = ["name", "count", "image", "text", *EXT2COLOR.keys()]

TEXT_FORMATS = {"txt", "caption", "wd", "json", "e621", "md", "tags"}
IMAGE_FORMATS = EXT2COLOR.keys() - TEXT_FORMATS 


def setup_logging(debug=False):
    if debug:
        logging.basicConfig(level=logging.DEBUG, filename='debug.log', filemode='w',
                            format='%(asctime)s - %(levelname)s - %(message)s')
    else:
        logging.basicConfig(level=logging.INFO)


def count_files(directory, debug=False):
    if debug:
        logging.debug(f"Scanning directory: {directory}")
    file_counts = Counter()
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.startswith('.'):
                if debug:
                    logging.debug(f"Skipping hidden file: {os.path.join(root, file)}")
                continue
            base_name, ext = os.path.splitext(file)
            if len(ext) > 1:
                if debug:
                    logging.debug(f"Processing file: {file}, Base: {base_name}, Extension: {ext}")
                ext = ext[1:].lower()
                file_counts[ext] += 1
                if ext in IMAGE_FORMATS:
                    file_counts["image"] += 1
                elif ext in TEXT_FORMATS:
                    file_counts["text"] += 1
            else:
                file_counts["no_extension"] += 1
                if debug:
                    logging.debug(f"File without extension: {os.path.join(root, file)}")
    return file_counts


def main():
    sort_key_name = "name"
    sort_reverse = False
    debug = False

    # Parse command-line arguments
    args = sys.argv[1:]
    if "--debug" in args:
        debug = True
        args.remove("--debug")
    
    # Set up logging based on debug flag
    setup_logging(debug)

    if debug:
        logging.debug(f"Debug mode enabled")
        logging.debug(f"Sort key name: {sort_key_name}, Sort reverse: {sort_reverse}")

    if len(args) > 0:
        sort_key_name = args[0]
    if sort_key_name.endswith("_r"):
        sort_reverse = True
        sort_key_name = sort_key_name[:-2]

    # Set up sort_key based on sort_key_name
    if sort_key_name == "name":
        sort_key = lambda x: x[0]
    elif sort_key_name == "count":
        sort_key = lambda x: x[1]
    elif sort_key_name in SORT_KEYS:
        sort_key = lambda x: x[2].get(sort_key_name, 0)
    else:
        print(_('Valid sort keys are {}').format(", ".join(f'"{k}"' for k in SORT_KEYS)))
        print(_('Append "_r" to reverse the sort order'))
        sys.exit(1)

    current_directory = os.getcwd()
    directories = (
        d
        for d in os.listdir(current_directory)
        if os.path.isdir(os.path.join(current_directory, d))
    )

    stats = []
    grand_total = Counter()
    for directory in directories:
        dir_path = os.path.join(current_directory, directory)
        counts = count_files(dir_path, debug=debug)
        total_files = sum(v for k, v in counts.items() if k in EXT2ORDER)
        stats.append((directory, total_files, counts))
        grand_total.update(counts)

    stats.sort(key=sort_key, reverse=sort_reverse)
    stats.append((None, sum(v for k, v in grand_total.items() if k in EXT2ORDER), grand_total))

    for directory, total_files, counts in stats:
        if total_files == 0:
            continue
        if directory is None:
            print(_('Grand Total:'))
        else:
            print(_('Directory: {}').format(directory))
            for txt_format in TEXT_FORMATS:
                if counts.get(txt_format, 0) == 0:
                    print(f"{RED}{_('❌ {txt_format} not found').format(txt_format=txt_format)}{RESET}")
        if counts.get("no_extension", 0) > 0:
            print(f"{ORANGE}⚠️ {_('Files without extensions found: {}').format(counts['no_extension'])}{RESET}")
        for ext, count in sorted(
            counts.items(), key=lambda x: EXT2ORDER.get(x[0], -1)
        ):
            if counts[ext] == 0 or ext not in EXT2COLOR:
                continue
            print(f"{EXT2COLOR[ext]}{ext} {_('files')}: {counts[ext]}{RESET}")
        tally_color = GREEN if total_files >= 200 else ORANGE
        print(f"{tally_color}{_('Total files: {}').format(total_files)}{RESET}")
        print()


if __name__ == "__main__":
    main()