Parse FASTA

parse_fasta.py

from functools import lru_cache
import os
import pandas as pd

@lru_cache(maxsize=16)
def read_fasta(fasta_fp):
    assert os.path.exists(fasta_fp)
    
    fasta = {}
    header = None
    seq = []
    
    with open(fasta_fp, "r") as handle:
        for line in handle:
            if line[0] == ">":
                if header is not None and len(seq) > 0:
                    fasta[header] = "".join(seq)
                header = line[1:].split(" ")[0].rstrip("\n")
                seq = []
            else:
                if len(line) > 1:
                    seq.append(line.rstrip("\n"))
                    
    if header is not None and len(seq) > 0:
        fasta[header] = "".join(seq)
        
    # Format as a DataFrame
    df = pd.DataFrame([
        {
            "header": header,
            "seq": seq,
            "len": len(seq)
        }
        for header, seq in fasta.items()
    ]).sort_values(by="len", ascending=False)
    
    # Set the row index by header
    df.set_index("header", inplace=True)
    
    return df