Return to Revisions

You could instead pass in a generator that yields SeqRecords. This avoids storing the entire list of sequences in memory and is usually the way to go assuming that you need to handle large files containing thousands of records. This example use a generator expression, which is a bit like a list comprehension, but instead of creating a list, it creates a generator:

import argparse
import contextlib
import pathlib
import sys

from Bio import SeqIO

def get_argument_parser():

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "fasta_file",
        type=pathlib.Path,
        metavar='FILE',
        help="Extract sequence records from FILE",
    )
    parser.add_argument(
        "ids_file",
        type=pathlib.Path,
        metavar='FILE',
        help="Sequence record IDs/names to extract, one per line",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=pathlib.Path,
        metavar='FILE',
        default="-",
        help="Write extracted sequences to FILE (default: stdout)",
    )
    parser.add_argument(
        "-f",
        "--force",
        action='store_true',
        help="Overwrite the output file if it exists",
    )

    return parser

def main():

    parser = get_argument_parser()
    args = parser.parse_args()

    with contextlib.ExitStack() as stack:

        ids_file = stack.enter_context(open(args.ids_file))
        fasta_file = stack.enter_context(open(args.fasta_file))

        if args.output != '-':
            out = stack.enter_context(open(args.output, 'w' if args.force else 'x'))
        else:
            out = sys.stdout

        record_ids = {line.rstrip('\n') for line in ids_file}

        fasta = SeqIO.parse(fasta_file, "fasta")
        records = (rec for rec in fasta if rec.id in record_ids)

        SeqIO.write(records, out, "fasta")

if __name__ == "__main__":
    main()