How to mock directory-based sharding in Python

import os
import hashlib
from collections import defaultdict
from pathlib import Path


def get_shard_for_key(key: str, num_shards: int) -> int:
    """Return a deterministic shard index (0..num_shards-1) for a key."""
    digest = hashlib.md5(key.encode('utf-8')).hexdigest()
    return int(digest, 16) % num_shards


def distribute_files_to_shards(source_dir: str, num_shards: int) -> dict:
    """Distribute files from a flat source directory into logical shards."""
    shard_buckets = defaultdict(list)
    source_path = Path(source_dir)

    if not source_path.is_dir():
        raise ValueError(f"Not a directory: {source_path}")

    for file_path in source_path.iterdir():
        if not file_path.is_file():
            continue
        filename = file_path.name
        shard_index = get_shard_for_key(filename, num_shards)
        shard_buckets[shard_index].append(filename)

    return dict(sorted(shard_buckets.items()))


if __name__ == "__main__":
    demo_dir = Path("./data")
    demo_dir.mkdir(exist_ok=True)

    for name in ("alice.csv", "bob.txt", "carol.log", "dave.json",
                 "eve.txt", "frank.csv", "grace.md", "heidi.log"):
        (demo_dir / name).write_text("sample content\n")

    shards = distribute_files_to_shards(str(demo_dir), num_shards=3)

    for shard_id, files in shards.items():
        print(f"Shard {shard_id}: {', '.join(files)}")