biyelunwen/99.scripts/miscs/make_bpp_seq.py

85 lines
2.7 KiB
Python

#!/usr/bin/env python
"""
Prepare input sequence file from fasta file in a directory for BPP.
"""
import argparse
import sys
from pathlib import Path
from Bio import SeqIO
def get_fasta_files(input_dir: Path, ext=".fasta") -> list[Path]:
"""
Get a list of fasta files in the specified directory with the given extension.
Args:
input_dir (Path): Path to the directory containing fasta files.
ext (str): Extension of the fasta files.
Returns:
list[Path]: List of Paths to the fasta files.
"""
files = list(input_dir.glob(f"*{ext}"))
if not files:
print(f"No fasta files with extension '{ext}' found in directory '{input_dir}'")
sys.exit(1)
return files
def transform_fasta_to_bpp(fasta_file: Path):
"""
Transform a fasta file into BPP sequence format.
Args:
fasta_file (Path): Path to the input fasta file.
Returns:
list[str]: Lines formatted for BPP input.
"""
num_ind = 0
output_lines = []
seq_len = []
for record in SeqIO.parse(fasta_file, "fasta"):
ind_id = record.id
num_ind += 1
output_lines.append(f"^{ind_id} {str(record.seq)}\n")
seq_len.append(len(record.seq))
if len(set(seq_len)) != 1:
print(f"Error: Sequences in file '{fasta_file}' have different lengths.")
sys.exit(1)
header = f"{num_ind} {seq_len[0]}\n"
return [header, "\n", *output_lines, "\n"]
def make_bpp_seq(input_dir: Path, output_file: Path, ext=".fasta"):
"""
Create a BPP formatted sequence file from fasta files in a directory.
Args:
input_dir (Path): Path to the directory containing fasta files.
output_file (Path): Path to the output BPP formatted file.
ext (str): Extension of the fasta files.
"""
fasta_files = get_fasta_files(input_dir, ext)
with open(output_file, "w") as out_f:
for fasta_file in fasta_files:
bpp_lines = transform_fasta_to_bpp(fasta_file)
out_f.writelines(bpp_lines)
print(f"BPP formatted sequence file written to '{output_file}'")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Create a BPP formatted sequence file from fasta files in a directory."
)
parser.add_argument(
"input_dir", type=Path, help="Directory containing fasta files."
)
parser.add_argument(
"output_file", type=Path, help="Output BPP formatted sequence file."
)
parser.add_argument(
"--ext",
type=str,
default=".fasta",
help="Extension of the fasta files (default: .fasta).",
)
args = parser.parse_args()
make_bpp_seq(args.input_dir, args.output_file, args.ext)