使用pickle模块更新学生信息时隐藏测试失败问题排查

任务背景 #

系统中已存在存储学生信息列表的student_info.pkl文件，需复制该文件为updated_info.pkl用于更新学生信息。要求编写函数add_student_info，实现以下逻辑：

• 将新学生信息加入updated_info.pkl的列表中
• 若学生ID已存在则覆盖该学生所有信息
• 若新信息中缺少ID则不执行任何操作

学生信息列表结构示例：

student_info = [
    {
        'id': 20200000,
        'name': 'Nguyen Van A',
        'score': {
            'math': 7.8,
            'english': 8.9,
            'physics': 9.0,
        }
    },
    {
        'id': 20200001,
        'name': 'Le Van B',
        'score': {
            'math': 9.8,
            'english': 8.7,
            'physics': 7.6,
        }
    }
]

测试用例：

new_student = {
        'id': 20200000,
        'name': 'Nguyen Van Anh',
        'score': {
            'math': 7.8,
            'english': 8.4,
            'physics': 8.0,
        }
    }
add_student_info(new_student)

原提交代码：

from shutil import copyfile
copyfile('student_info.pkl', 'updated_info.pkl')
#my code starts from here
def add_student_info(ne):
    with open('updated_info.pkl', 'rb') as f:
        a = pickle.load(f)
    for i in range(len(a)):
        if a[i]['id'] == ne['id']:
            a[i] = ne
            continue
    with open('updated_info.pkl', 'wb') as f:
        pickle.dump(a,f)

原代码的失败场景 #

• 未处理缺少id的新学生数据：直接访问ne['id']会触发KeyError，不符合“缺少ID则不执行任何操作”的要求。
• 未添加新学生（ID不存在时）：仅处理了ID存在时的覆盖逻辑，完全忽略了ID不存在时需要将新学生追加到列表的需求。
• 循环效率低下且逻辑冗余：通过range(len(a))遍历列表，找到匹配项后用continue继续遍历，没有及时跳出循环，浪费资源。
• 缺失pickle模块导入：代码中使用了pickle的方法但未导入模块，运行时会报错。
• 文件复制逻辑位置错误：copyfile放在函数外，若多次调用add_student_info，每次都会从原始文件覆盖updated_info.pkl，导致之前的更新全部丢失。

优化后的代码 #

import pickle
from shutil import copyfile
import os

def add_student_info(new_student):
    # 检查新学生数据是否包含id，不包含则直接返回
    if 'id' not in new_student:
        return
    
    # 仅当目标文件不存在时复制原始文件，避免覆盖已有更新
    if not os.path.exists('updated_info.pkl'):
        copyfile('student_info.pkl', 'updated_info.pkl')
    
    # 读取现有学生信息列表
    with open('updated_info.pkl', 'rb') as f:
        student_list = pickle.load(f)
    
    # 遍历查找相同ID的学生，找到后覆盖并标记更新状态
    updated = False
    for idx, student in enumerate(student_list):
        if student['id'] == new_student['id']:
            student_list[idx] = new_student
            updated = True
            break  # 找到匹配项后立即跳出循环，提升效率
    
    # 若未找到相同ID的学生，将新学生追加到列表
    if not updated:
        student_list.append(new_student)
    
    # 将更新后的列表写入文件
    with open('updated_info.pkl', 'wb') as f:
        pickle.dump(student_list, f)

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