folding-studio/folding_studio/utils/input_validation.py

"""Helper methods for validating cli entries."""

from enum import Enum
from pathlib import Path

import typer

from folding_studio.utils.path_helpers import extract_files


class SupportedCustomTemplateFile(str, Enum):
    """Supported custom template file extensions."""

    CIF = ".cif"


class SupportedCustomMSAFile(str, Enum):
    """Supported custom msa file extensions."""

    STO = ".sto"
    A3M = ".a3m"


class SupportedTemplateMaskFile(str, Enum):
    """Supported custom template file extensions."""

    JSON = ".json"


def extract_and_validate_custom_templates(
    paths: list[Path],
) -> list[Path]:
    """Extract and validate a list of path to use as custom templates.

    Before checking the files, if a directory path is in the list, its files are extracted.

    Args:
        paths (list[Path]): List of paths.

    Raises:
        typer.BadParameter: If a file is not supported.

    Returns:
        list[Path]: Extracted and validated custom template files path.
    """
    extracted_paths = extract_files(paths)
    supported_files = [item.value for item in SupportedCustomTemplateFile]
    for path in extracted_paths:
        if path.suffix not in supported_files:
            raise typer.BadParameter(
                f"The file '{path}' is not supported. "
                f"Only {supported_files} are supported for custom templates."
            )
    return extracted_paths


def extract_and_validate_custom_msas(paths: list[Path]) -> list[Path]:
    """Extract and validate a list of path to use as custom msas.

    Before checking the files, if a directory path is in the list, its files are extracted.

    Args:
        paths (list[Path]): List of paths.

    Raises:
        typer.BadParameter: If a file is not supported.

    Returns:
        list[Path]: Extracted and validated custom msas files path.
    """
    extracted_paths = extract_files(paths)
    supported_files = [item.value for item in SupportedCustomMSAFile]
    for path in extracted_paths:
        if path.suffix not in supported_files:
            raise typer.BadParameter(
                f"The file '{path}' is not supported. "
                f"Only {supported_files} are supported for custom msas."
            )
    return extracted_paths


def validate_initial_guess(
    initial_guess_file: Path | None,
) -> Path:
    """Extract and validate the file used for initial guess.

    Args:
        path (Path | None): Initial guess file parameter.

    Raises:
        typer.BadParameter: If a file is not supported.

    Returns:
        Path | None: Validated initial guess file path or None.
    """
    if not initial_guess_file:
        return None
    supported_files = [
        item.value for item in SupportedCustomTemplateFile
    ]  # initial guess files are equivalent to template files
    if initial_guess_file.suffix not in supported_files:
        raise typer.BadParameter(
            f"The file '{initial_guess_file}' is not supported. "
            f"Only {supported_files} are supported for initial guess."
        )
    return initial_guess_file


def validate_template_mask(
    template_mask_file: Path | None,
) -> Path:
    """Extract and validate the file used for template masks.g.

    Args:
        path (Path | None): template masks.g file parameter.

    Raises:
        typer.BadParameter: If a file is not supported.

    Returns:
        Path | None: Validated template masks.g file path or None.
    """
    if not template_mask_file:
        return None
    supported_files = [
        item.value for item in SupportedTemplateMaskFile
    ]  # template masks.g files are equivalent to template files
    if template_mask_file.suffix not in supported_files:
        raise typer.BadParameter(
            f"The file '{template_mask_file}' is not supported. "
            f"Only {supported_files} are supported for template masks.g."
        )
    return template_mask_file


def validate_list_initial_guess(ig_files: list[Path]) -> list[Path]:
    """Extract and validate a list of initial guess files.

    Args:
        ig_files (list[Path]): List of initial guess files path.

    Returns:
        list[Path]: Extracted and validated initial guess files path.

    """
    extracted_paths = extract_files(ig_files)
    initial_guess_files = []
    for path in extracted_paths:
        validated_ig_path = validate_initial_guess(initial_guess_file=path)
        if validated_ig_path:
            initial_guess_files.append(validated_ig_path)

    return initial_guess_files


def validate_list_template_mask(tm_files: list[Path]) -> list[Path]:
    """Extract and validate a list of template mask

    Args:
        tm_files (list[Path]): List of template mask files path.

    Returns:
        list[Path]: Extracted and validated template mask files path.

    """
    extracted_paths = extract_files(tm_files)
    template_mask_files = []
    for path in extracted_paths:
        validated_tm_path = validate_template_mask(template_mask_file=path)
        if validated_tm_path:
            template_mask_files.append(validated_tm_path)

    return template_mask_files