# -*- coding: utf-8 -*-
import os
import sys
import shutil
from pprint import pprint
def extract_fasta(file1, match_list):
"""
从file1里检索match_list的内容，并写入new_file
:param file1: 输入文件的路径
:param match_list: list 格式的检索字段
:return: new_file: 新文件的路径
"""
if not os.path.isabs(file1):
file1 = os.path.abspath(file1)
part = file1.rpartition('.')
new_file = part[0] + '_1_taxa.' + part[2]
with open(file1) as f1, open(new_file, 'w') as f2:
content = f1.readlines()
lines_num = len(content)
for match_str in match_list:
for line_num in range(lines_num):
if match_str in content[line_num]:
f2.write(content[line_num])
f2.write(content[line_num+1])
break
return new_file
def batch_extract_fasta(path, match_list, file_type='.fas'):
"""
读取 path 路径下的所有 fasta 文件，并提取出现在 match_list 中的字符串至当前路径下的 new 文件夹中
:param path: fasta 文件所在的路径
:param match_list: 需要匹配的字符串列表
:param file_type: None
:return:
"""
if not os.path.isabs(path):
path = os.path.abspath(path)
result_path = os.path.join(path, 'new')
if not os.path.exists(result_path):
os.makedirs(result_path)
file_list = [os.path.join(path, file) for file in os.listdir(path) if file.endswith(file_type)]
for file in file_list:
new_file = extract_fasta(file, match_list)
# 将结果移动至 new 文件夹下
part = new_file.rpartition(os.sep)
result_file = os.path.join(part[0], 'new', part[2])
shutil.move(new_file, result_file)
# cmd 直接传参调用
if __name__ == '__main__':
path = sys.argv[1]
file = sys.argv[2]
with open(file) as f:
content_list = [line.strip() for line in f.readlines() if len(line) > 1]
batch_extract_fasta(path, content_list)
# 内部调试代码块
# if __name__ == '__main__':
# path = 'data'
# file = 'extracted.txt'
# with open(file) as f:
# content_list = [line.strip() for line in f.readlines() if len(line) > 1]
# batch_extract_fasta(path, content_list)
评论2