批量修改基因文件里的基因名

前面提到了可以根据树的拓扑结构生成一系列与物种树一致的基因树；但其实也可以用同一个物种树，然后把OG基因里的基因名替换成物种名。
同样也需要一个genelist文件（这个文件生成参考前面的），以及一个输入的OG基因文件目录。

#!/usr/bin/python
# -*- coding: utf-8 -*-
# conda activate python3
"""
    作者：徐诗芬
    内容：1. 打开single_copy_gene_list.txt文件，读取第一行生成一个species list；
         2. 读取剩下的行，读取OG的id以及生成该id的基因列表gene list,
         并且将species和gene列表并列起来生成一个字典；
         3.根据OG的id打开Single_copy_gene里的id文件，读取gene名，循环gene list，根据基因名替换成物种，
         创建一个与OG_id相同的文件直接写入替换好的序列文件，用writelines可以直接写一个列表。

    日期：2022.1.11
"""
import sys
import os
import subprocess
import re
from itertools import islice


def usage():
    print('Usage: python change_gene.py [single_copy_gene_list.txt] [Input_dir] [Output_dir]')

def main():

    name_list = open('single_copy_gene_list.txt', 'rt')  # sys.argv[1]
    input_dir = 'data\Single_copy_gene' #sys.argv[2]  # 输入文件夹，Single_copy_gene,
    out_dir = 'data\Single_copy_gene_species' #sys.argv[3]   # 运行前先创建一个，可以命名为Single_copy_gene_species
    # subprocess.call(['mkdir', out_dir])  # linux系统自动创建output文件夹

    species_list = []
    name_list.seek(0, 0)                   # 加一个指针文件，运行完下面的第一行，后面的readline就可以直接从下一行开始
    for head in islice(name_list, 0, 1):   # 输出第一行title
    # head = name_list.readline()           # 也可以这样输出title
        species_list = head.strip().split("\t")[1:]

    # line = name_list.readline()
    for line in name_list:                          # 由于上面指定了文件开始位置，这里其实就是跳过第一行读取剩下的行
        OG_id = line.strip().split("\t")[0]         # 读取OG的id
        gene_list = line.strip().split("\t")[1:]    # 读取该行所包含的所有物种基因名
        pattern_dict = dict(zip(gene_list, species_list))  # 生成物种与基因名一一对应的字典
        OG_list = os.listdir(input_dir)              # 生成一个输入文件目录里存在的文件列表
        infile = f'{OG_id}.fa.mafft.fa'
        if infile in OG_list:                        # 判断infile是否在input_dir里面
            fas_file = os.path.join(input_dir, infile)
            pref = infile.split('.')[0]
            with open(fas_file, 'r') as f:
                new_lines = []
                for fa_line in f:
                    if fa_line.startswith('>'):
                        name = fa_line.strip()[1:]
                        if name in list(pattern_dict.keys()):
                            new_line = re.sub(name, pattern_dict[name], fa_line)  # 把基因名替换成该行里面的物种名
                            new_lines.append(new_line)
                    else:
                        new_lines.append(fa_line)

                with open(f'{out_dir}/{pref}.mafft.fa', 'w') as out:
                    out.writelines(new_lines)


try:
    main()
except IndexError:
    usage()