|
|
|
|
|
import os |
|
import sys |
|
from collections import Counter |
|
|
|
|
|
|
|
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, |
|
"txt": BLUE, |
|
} |
|
EXT2ORDER = {ext: i for i, ext in enumerate(EXT2COLOR.keys())} |
|
SORT_KEYS = ["name", "count", "image", "text", *EXT2COLOR.keys()] |
|
|
|
TEXT_FORMATS = {"txt", "caption"} |
|
IMAGE_FORMATS = EXT2COLOR.keys() - TEXT_FORMATS |
|
|
|
|
|
def count_files(directory): |
|
file_counts = Counter() |
|
for root, dirs, files in os.walk(directory): |
|
for file in files: |
|
base_name, ext = os.path.splitext(file) |
|
if len(ext) > 1: |
|
ext = ext[1:] |
|
file_counts[ext] += 1 |
|
if ext in IMAGE_FORMATS: |
|
file_counts["image"] += 1 |
|
elif ext in TEXT_FORMATS: |
|
file_counts["text"] += 1 |
|
|
|
return file_counts |
|
|
|
|
|
def main(): |
|
sort_key_name = "name" |
|
sort_reverse = False |
|
if len(sys.argv) > 1: |
|
sort_key_name = sys.argv[1] |
|
if sort_key_name.endswith("_r"): |
|
sort_reverse = True |
|
sort_key_name = sort_key_name[:-2] |
|
|
|
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(f'Valid short key are {", ".join(f'"{k}"' for k in SORT_KEYS)}') |
|
print('Prepending "_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) |
|
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(f'Grand Total: ') |
|
print(f"Directory: {directory}") |
|
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: {total_files}{RESET} ({counts['image']} images, {counts['text']} texts)" |
|
) |
|
print() |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
