work-helpers/work_helpers/sensospot_rename.py

""" rename sensospot images """

import click
import pathlib
import shutil


RENAME_MAP_NAME = "rename_map.txt"
IMAGE_SUFFIXES = {".tif", ".jpg"}


def get_map_file_path(directory:str) -> pathlib.Path:
    return pathlib.Path(directory) / RENAME_MAP_NAME


def list_images(directory:str) -> list[pathlib.Path]:
    parent = pathlib.Path(directory)
    non_hidden = (i for i in parent.iterdir() if not i.name.startswith("."))
    return [i for i in non_hidden if i.suffix.lower() in IMAGE_SUFFIXES]


def get_image_parts(image_path:pathlib.Path) -> tuple[str, str, str]:
    return image_path.stem.rsplit("_", maxsplit=2)


def get_unique_parts(images:list[pathlib.Path]) -> list[list[str], list[str], list[str]]: 
    parts = [get_image_parts(p) for p in images]
    return [sorted(set(items)) for items in zip(*parts)]


def write_rename_map(map_file:pathlib.Path) -> None:
    images = list_images(map_file.parent)
    parts = get_unique_parts(images)
    headers = ["stems", "wells", "exposures"]
    
    lines = []
    for header, items in zip(headers, parts):
        lines.append(f"# {header}")
        lines.extend([f"{i}\t" for i in items])
        lines.append("")
    lines.append("")
    
    map_file.write_text("\n".join(lines))


def _parse_rename_map(content:str) -> tuple[str, str]:
    lines = [line.strip() for line in content.splitlines()]
    for i, line in enumerate(lines, start=1):
        if not line or line.startswith("#"):
            continue
        if not "\t" in line:
            raise ValueError(f"No tab in line {i}: '{line}'")
        yield line.split("\t", maxsplit=1)


def read_rename_map(map_file:pathlib.Path) -> dict[str, str]:
    content = map_file.read_text()
    return {k: v for k, v in _parse_rename_map(content)}


def prepare_rename(images:list[pathlib.Path], rename_map:dict[str, str], sep="_") -> list[tuple[pathlib.Path, pathlib.Path]]:
    for path in images:
        renamed_parts = [rename_map[p] for p in get_image_parts(path)]
        yield path, path.with_stem(sep.join(renamed_parts))


@click.command()
@click.argument("directory", type=click.Path(exists=True, file_okay=False, dir_okay=True))
def sensospot_rename(directory):
    images = list_images(directory)
    if not images:
        raise click.UsageError(f"No images found in '{directory}'")

    map_file = get_map_file_path(directory)
    if not map_file.is_file():
        write_rename_map(map_file)
        click.echo(f"Prepared rename map at '{map_file!s}'")
        click.echo(f"Rerun the command after editing the file.")
    else:
        click.echo(f"Reading rename map at '{map_file!s}'")
        try:
            rename_map = read_rename_map(map_file)
        except ValueError as e:
            raise click.UsageError(str(e))
        try:
            prepared = list(prepare_rename(images, rename_map))
        except Exception:
            raise click.UsageError("Could not rename images. Please check the image directory and rename map file.")

        for src, dst in prepared:
            click.echo(f"renaming: {src} -> {dst}")
            shutil.move(src, dst)