File size: 3,099 Bytes
51b7ba7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
#!/usr/bin/env python3
import os
import sys
from collections import Counter
# 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,
"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()
|