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 sys
from Bio import SeqIO
def get_argument_parser():
parser = argparse.ArgumentParser()
parser.add_argument(
"fasta_file",
type=str,
metavar='FILE',
help="Extract sequence records from FILE",
)
parser.add_argument(
"ids_file",
type=str,
metavar='FILE',
help="Sequence record IDs/names to extract, one per line",
)
parser.add_argument(
"-o",
"--output",
type=str,
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()
