41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
|
from __future__ import annotations
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Iterable, Set
|
||
|
|
|
||
|
|
|
||
|
|
def read_id_set(path: str) -> Set[str]:
|
||
|
|
s = set()
|
||
|
|
with open(path, "r", encoding="utf-8", errors="ignore") as f:
|
||
|
|
for line in f:
|
||
|
|
x = line.strip()
|
||
|
|
if x:
|
||
|
|
s.add(x)
|
||
|
|
return s
|
||
|
|
|
||
|
|
|
||
|
|
def extract_fasta_by_ids(in_fasta: str, ids: Set[str], out_fasta: str) -> int:
|
||
|
|
"""
|
||
|
|
依据 fasta header 的第一个 token(>ID 后空格前) 提取序列
|
||
|
|
返回写出的序列条数
|
||
|
|
"""
|
||
|
|
in_path = Path(in_fasta)
|
||
|
|
out_path = Path(out_fasta)
|
||
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
keep = False
|
||
|
|
n = 0
|
||
|
|
|
||
|
|
with in_path.open("r", encoding="utf-8", errors="ignore") as fin, \
|
||
|
|
out_path.open("w", encoding="utf-8") as fout:
|
||
|
|
for line in fin:
|
||
|
|
if line.startswith(">"):
|
||
|
|
seq_id = line[1:].strip().split()[0]
|
||
|
|
keep = seq_id in ids
|
||
|
|
if keep:
|
||
|
|
fout.write(line)
|
||
|
|
n += 1
|
||
|
|
else:
|
||
|
|
if keep:
|
||
|
|
fout.write(line)
|
||
|
|
return n
|