splitter.py

from pathlib import Path
import re
from collections import Counter


def sanitize_filename(name: str) -> str:
    """Make a safe filename from a FASTA header token."""
    name = name.strip()
    name = re.sub(r"[^\w\-.]", "_", name)
    return name or "sequence"



def split_fasta_per_sequence(input_text: str, output_dir: str | Path):
    """
    Split FASTA text so each sequence is written to its own file,
    using the first token from the header as the filename.

    Returns:
        tuple[list[Path], int]: created file paths and sequence count
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    header = None
    seq_lines = []
    seq_count = 0
    created_files = []
    used_names = Counter()

    for raw_line in input_text.splitlines():
        line = raw_line.strip()
        if not line:
            continue

        if line.startswith(">"):
            if header is not None:
                seq_name = sanitize_filename(header[1:].split()[0])
                used_names[seq_name] += 1
                if used_names[seq_name] > 1:
                    seq_name = f"{seq_name}_{used_names[seq_name]}"

                out_path = output_dir / f"{seq_name}.fasta"
                with open(out_path, "w", encoding="utf-8") as out_f:
                    out_f.write(header + "\n")
                    out_f.write("\n".join(seq_lines) + "\n")
                created_files.append(out_path)
                seq_count += 1

            header = line
            seq_lines = []
        else:
            seq_lines.append(line)

    if header is not None:
        seq_name = sanitize_filename(header[1:].split()[0])
        used_names[seq_name] += 1
        if used_names[seq_name] > 1:
            seq_name = f"{seq_name}_{used_names[seq_name]}"

        out_path = output_dir / f"{seq_name}.fasta"
        with open(out_path, "w", encoding="utf-8") as out_f:
            out_f.write(header + "\n")
            out_f.write("\n".join(seq_lines) + "\n")
        created_files.append(out_path)
        seq_count += 1

    return created_files, seq_count