set-soft-AudioSeparation/source/db/hash_dir.py

#!/usr/bin/env python3
# Copyright (c) 2025 Salvador E. Tropea
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
# License: GPLv3
# Project: ComfyUI-AudioSeparation
#
# Tool to get hashes for all the files in a dir, using a cache
# Can be used as an standalone tool for testing
# Code by Gemini 2.5 Pro
import os
import csv
import logging
import sys

from .hash import get_hash
from ..utils.misc import NODES_NAME, debugl

# Set up the logger as specified
logger = logging.getLogger(f"{NODES_NAME}.hash_dir")

# Constants for clarity
MIN_FILE_SIZE_MB = 10
MIN_FILE_SIZE_BYTES = MIN_FILE_SIZE_MB * 1024 * 1024
CATALOG_FILENAME = ".catalog.csv"


def hash_dir(directory_path: str) -> dict:
    """
    Computes hashes for large files in a directory, using a local cache
    to avoid re-computation.

    Args:
        directory_path: The path to the directory to scan.

    Returns:
        A dictionary mapping {file_hash: file_name} for all files in the
        directory that are 10MB or larger.
    """
    logger.debug(f"Starting hash process for directory: '{directory_path}'")

    # 1. Ensure the target directory exists.
    try:
        os.makedirs(directory_path, exist_ok=True)
    except OSError as e:
        logger.error(f"FATAL: Could not create directory '{directory_path}'. Error: {e}")
        # This is a fatal error, so we raise the exception to stop execution.
        raise

    catalog_path = os.path.join(directory_path, CATALOG_FILENAME)
    cached_data = {}

    # 2. Load existing cache from ".catalog.csv" if it exists.
    try:
        with open(catalog_path, 'r', newline='') as f:
            reader = csv.reader(f)
            # Skip header
            next(reader, None)
            for row in reader:
                if len(row) == 3:
                    filename, file_hash, timestamp = row
                    cached_data[filename] = (file_hash, float(timestamp))
            logger.debug(f"Successfully loaded {len(cached_data)} entries from cache: {catalog_path}")
    except FileNotFoundError:
        logger.debug(f"Cache file '{catalog_path}' not found.")
    except Exception as e:
        logger.warning(f"Could not read cache file '{catalog_path}'. Proceeding without cache. Error: {e}")

    final_hashes = {}
    updated_cache = {}
    has_updates = False

    # 3. Iterate through files in the directory.
    for filename in os.listdir(directory_path):
        file_path = os.path.join(directory_path, filename)

        # Skip subdirectories and the cache file itself
        if not os.path.isfile(file_path) or filename == CATALOG_FILENAME:
            continue

        # 4. Filter by file size.
        try:
            file_size = os.path.getsize(file_path)
            if file_size < MIN_FILE_SIZE_BYTES:
                logger.debug(f"Skipping small file: '{filename}' ({file_size / 1024**2:.2f}MB)")
                continue
        except OSError as e:
            logger.warning(f"Could not get size of file '{filename}'. Skipping. Error: {e}")
            continue

        current_mtime = os.path.getmtime(file_path)
        file_hash = None

        # 5. Check cache for a valid, up-to-date entry.
        if filename in cached_data:
            cached_hash, cached_mtime = cached_data[filename]
            if current_mtime == cached_mtime:
                debugl(logger, 2, f"Cache hit for '{filename}'. Using stored hash.")
                file_hash = cached_hash
            else:
                logger.debug(f"File '{filename}' has been modified. Re-calculating hash.")

        # 6. If no valid cache entry, compute the hash.
        if file_hash is None:
            logger.debug(f"Computing hash for '{filename}'...")
            try:
                file_hash = get_hash(file_path)
                has_updates = True
            except IOError as e:
                logger.warning(f"Could not read file '{filename}' to compute hash. Skipping. Error: {e}")
                continue

        # Add to the results and prepare for caching
        final_hashes[file_hash] = file_path
        updated_cache[filename] = (file_hash, current_mtime)

    # 7. Write the updated cache back to disk.
    if has_updates:
        try:
            with open(catalog_path, 'w', newline='') as f:
                writer = csv.writer(f)
                writer.writerow(['Filename', 'Hash', 'Time stamp'])
                for filename, (file_hash, mtime) in updated_cache.items():
                    writer.writerow([filename, file_hash, mtime])
            logger.debug(f"Successfully wrote {len(updated_cache)} entries to cache file '{catalog_path}'.")
        except IOError as e:
            # This is a non-fatal warning as per requirements.
            logger.warning(f"Failed to write to cache file '{catalog_path}'. Hashes were computed but not saved. Error: {e}")

    logger.debug(f"Hash process finished. Found {len(final_hashes)} valid files.")
    return final_hashes


# ==============================================================================
# Command-Line Tool for Testing and Operation
# python -m source.db.hash_dir
# ==============================================================================
if __name__ == "__main__":
    # Local imports to avoid top-level pollution
    import argparse
    from pathlib import Path
    import pprint

    # --- Argument Parsing ---
    parser = argparse.ArgumentParser(
        description="Calculate and cache hashes for large files in a directory.",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        'directory',
        nargs='?',  # Makes the argument optional; we'll validate it manually.
        help='The directory to scan. Required unless --test is used.'
    )
    parser.add_argument(
        '--test',
        action='store_true',  # This is a flag; if present, args.test will be True.
        help='Run in test mode. Creates a temporary directory with dummy files and runs a test sequence.'
    )
    args = parser.parse_args()

    # --- Setup Logging ---
    # Enable debug logging for command-line execution to provide useful feedback.
    logging.basicConfig(
        level=logging.DEBUG,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

    # --- Main Logic ---
    try:
        if args.test:
            # --- TEST MODE ---
            # This mode ignores the 'directory' argument and uses a fixed test path.

            # Helper function only needed for test mode
            def create_test_file(path: Path, size_in_bytes: int):
                """Helper to create a dummy file of a specific size."""
                path.parent.mkdir(parents=True, exist_ok=True)
                with open(path, "wb") as f:
                    f.write(os.urandom(size_in_bytes))
                logger.info(f"Created test file: {path} ({size_in_bytes / 1024**2:.2f}MB)")

            test_dir = Path("./temp_hash_dir_test_cli")
            print("-" * 60)
            print("Running in TEST mode.")
            print(f"Test directory: '{test_dir.resolve()}'")
            print("-" * 60)

            # Create test files
            create_test_file(test_dir / "large_file_1.bin", 11 * 1024 * 1024)  # > 10MB
            create_test_file(test_dir / "large_file_2.bin", 15 * 1024 * 1024)  # > 10MB
            create_test_file(test_dir / "small_file.txt", 1 * 1024 * 1024)   # < 10MB
            create_test_file(test_dir / "edge_case_file.bin", 10 * 1024 * 1024)  # exactly 10MB

            # Run the test sequence
            print("\n>>> FIRST RUN: Computing all hashes...")
            result_dict = hash_dir(str(test_dir))
            pprint.pprint(result_dict)

            print("\n>>> SECOND RUN: Should use cache for all files...")
            result_dict_cached = hash_dir(str(test_dir))
            pprint.pprint(result_dict_cached)

            print("\n>>> THIRD RUN: Modifying a file and re-running to test cache invalidation...")
            create_test_file(test_dir / "large_file_1.bin", 12 * 1024 * 1024)
            result_dict_modified = hash_dir(str(test_dir))
            pprint.pprint(result_dict_modified)

        else:
            # --- NORMAL OPERATION MODE ---

            # In normal mode, the 'directory' argument is required.
            if not args.directory:
                parser.error("The 'directory' argument is required when not using --test.")

            target_directory = args.directory
            print("-" * 60)
            print(f"Running in NORMAL mode on directory: '{target_directory}'")
            print("-" * 60)

            # Just run the function once on the specified directory
            result_dict = hash_dir(target_directory)

            print("\n>>> Resulting Hashes:")
            pprint.pprint(result_dict)

    except Exception as e:
        logger.critical(f"A critical error occurred during execution: {e}")
        sys.exit(1)