movement-sample-data/update_hashes.py

"""Compute sha256 hashes for files and update metadata.yaml."""

from pathlib import Path

import yaml
from pooch.hashes import file_hash


def _update_file_hash(
    file_path: Path,
    metadata_dict: dict,
    hash_key: str,
) -> None:
    """Update hash for a single file if the file exists and hash differs.

    Parameters
    ----------
    file_path : Path
        Path to the file to check and update hash for.
    metadata_dict : dict
        The metadata dictionary to update.
    hash_key : str
        The key in metadata_dict where the hash should be stored.

    """
    existing_hash = metadata_dict[hash_key]
    if file_path.exists():
        computed_hash = file_hash(file_path.as_posix())
        if existing_hash != computed_hash or existing_hash is None:
            metadata_dict[hash_key] = computed_hash
            print(
                f"Updated hash for {file_path.name}: "
                f"{existing_hash} -> {computed_hash}"
            )
    else:
        print(f"WARNING: Could not find {file_path}.")


def _update_dataset_hashes(sample_metadata: dict, root_dir: Path) -> None:
    """Update all hashes (dataset, frame, video) for a single dataset.

    Parameters
    ----------
    sample_metadata : dict
        Metadata dictionary for a single sample dataset.
    root_dir : Path
        Root of the sample data directory.

    """
    filename = list(sample_metadata.keys())[0]
    sample_metadata = sample_metadata[filename]
    ds_type = sample_metadata["type"]

    # Update main dataset hash
    dataset_path = root_dir / ds_type / filename
    _update_file_hash(
        file_path=dataset_path,
        metadata_dict=sample_metadata,
        hash_key="sha256sum",
    )
    
    # Update frame hash if frame file exists
    frame_file_name = sample_metadata["frame"]["file_name"]
    if frame_file_name is not None:
        frame_path = root_dir / "frames" / frame_file_name
        _update_file_hash(
            file_path=frame_path,
            metadata_dict=sample_metadata["frame"],
            hash_key="sha256sum",
        )
    
    # Update video hash if video file exists
    video_file_name = sample_metadata["video"]["file_name"]
    if video_file_name is not None:
        video_path = root_dir / "videos" / video_file_name
        _update_file_hash(
            file_path=video_path,
            metadata_dict=sample_metadata["video"],
            hash_key="sha256sum",
        )


def update_hashes_in_metadata(metadata_path: Path):
    """Update metadata.yaml with computed SHA256 hashes.

    Parameters
    ----------
    metadata_path : Path
        Path to the metadata.yaml file, which should be in the root
        of the sample data directory.

    """
    # Load existing metadata
    with open(metadata_path) as f:
        metadata = yaml.safe_load(f)

    # Get folder where metadata.yaml is located
    root_dir = metadata_path.parent

    # Update hashes for each dataset
    for sample_name in metadata:
        sample_metadata = {sample_name: metadata[sample_name]}
        _update_dataset_hashes(sample_metadata, root_dir)

    # Save updated metadata
    with open(metadata_path, 'w') as f:
        yaml.dump(
            metadata,
            f,
            default_flow_style=False,
            sort_keys=False,
        )

    print("Metadata updated successfully!")


if __name__ == "__main__":
    metadata_path = Path.cwd() / "metadata.yaml"
    update_hashes_in_metadata(metadata_path)