按样本ID拆分变异文件时如何避免重复写入VCF表头？

问题根源 #

原代码在处理每一行数据时，都会打开对应样本的文件并追加写入表头，导致同一个样本的每一条记录前都重复出现表头。同时还存在一个小错误：输入文件的CHROM列已经带chr前缀，代码里又额外拼接了一次，会生成chrchr1这种无效的染色体名。

修正方案 #

通过两点核心改动解决问题：

• 用集合记录已写入表头的样本，确保每个样本仅写入一次表头
• 复用文件句柄，减少重复打开/关闭文件的开销，同时避免资源泄漏

修正后的代码 #

def create_vcf(variant_file):
    # 定义符合预期的输出表头
    header = '#CHROM\tPOS\tREF\tALT\n'
    # 记录已写入表头的样本
    written_samples = set()
    # 保存每个样本的文件句柄，避免重复打开
    file_handles = {}
    
    with open(variant_file, "r") as f:
        next(f)  # 跳过输入文件的原始表头
        for line in f:
            line = line.strip()
            if not line:
                continue  # 跳过空行
            parts = line.split()
            sample_name = parts[0]
            chrom = parts[1]  # 直接使用输入文件的CHROM值，移除多余的chr拼接
            pos = parts[2]
            ref = parts[3]
            alt = parts[4]
            
            # 若样本文件未打开，先创建目录并打开文件
            if sample_name not in file_handles:
                file_path = f"{DIR_NAME}/output/{sample_name}/{sample_name}.vcf"
                # 自动创建不存在的目录
                import os
                os.makedirs(os.path.dirname(file_path), exist_ok=True)
                fh = open(file_path, "w")
                file_handles[sample_name] = fh
                # 写入表头并标记已完成
                fh.write(header)
                written_samples.add(sample_name)
            
            # 写入当前行的变异数据
            out_line = "\t".join([chrom, pos, ref, alt]) + "\n"
            file_handles[sample_name].write(out_line)
    
    # 处理完所有数据后统一关闭文件
    for fh in file_handles.values():
        fh.close()

关键改动说明 #

1. 表头写入控制：通过written_samples集合跟踪已写入表头的样本，仅在样本第一次被处理时写入表头。
2. 文件句柄复用：用file_handles字典保存每个样本的文件对象，避免每次处理一行都重复打开/关闭文件，提升处理效率。
3. 修正染色体名错误：移除原代码中多余的chr拼接，直接使用输入文件中的CHROM值。
4. 自动创建目录：新增os.makedirs确保输出目录存在，避免因目录不存在导致的写入报错（若已提前创建目录，可移除该部分代码）。

内容的提问来源于stack exchange，提问作者eb0906