python pickle文件正反向查找元素跳过问题的解决方案咨询

代码按顺序调用goto_index('1')、goto_index('2')、goto_index('3')读取pickle序列化文件时，预期输出hello world, good morning, 69 420，实际返回hello world, good morning, ending !，已定位根因为goto_index每次调用都会将内部计数obj_cnt重置为-1。

问题的具体触发逻辑：

• 第一次调用取索引1（对应0基计数0），从文件头开始读，计数从-1加到0匹配成功，seek回第一个对象的起始位置后返回，此时文件指针停在第一个对象开头
• 第二次调用取索引2（对应0基计数1），计数重置为-1，从第一个对象开头读，读完第一个对象计数为0不匹配，继续读第二个对象计数为1匹配成功，seek回第二个对象开头返回
• 第三次调用取索引3（对应0基计数2），计数再次重置为-1，从第二个对象开头读，读完第二个对象计数为0、第三个对象计数为1、第四个对象计数为2刚好匹配，所以返回了第四个对象ending !

提供两种适配不同场景的修复方案：

方案1：新增计数偏移参数（适配按升序读取、避免重复seek的高性能场景） #

修改goto_index函数，新增可选的start_cnt参数传入当前文件位置已读的对象数量，同时返回值追加当前匹配的索引，方便后续调用传参：

import pickle,os

def goto_index(idx_str,src,dest=False, start_cnt=-1) :
    '''
    新增参数：start_cnt  当前文件指针位置之前已经存在的对象数量，默认-1表示从文件头开始计数
    返回值修改为 (目标值, 匹配到的顶层对象索引)
    '''
    index = [int(subidx)-1 for subidx in idx_str.split(',')]
    val = None
    obj_cnt = start_cnt                             # 改用传入的初始计数

    try :
        while True :                         
            obj = pickle.load(src)
            obj_cnt += 1
            if obj_cnt == index[0] :
                val = obj
                for subidx in index[1::] :
                    val = val[subidx]        
                src.seek(-len(pickle.dumps(obj)),os.SEEK_CUR) 
                return val, obj_cnt                 # 返回值追加当前计数
            elif dest : pickle.dump(obj,dest)
    except (EOFError,IndexError) : raise

def add_elements(f) :
    pickle.dump('hello world',f)
    pickle.dump('good morning',f)
    pickle.dump('69 420',f)
    pickle.dump('ending !',f)

def get_elements(f) :
    elements = []
    current_cnt = -1
    for idx_str in ('1','2','3') : 
        val, current_cnt = goto_index(idx_str,f, start_cnt=current_cnt)
        elements.append(val)
        # 读取完成后指针停在当前对象开头，下一次读取前先跳过当前对象
        pickle.load(f)
        current_cnt +=1
    return elements

with open("tmp","wb+") as tmp :
    add_elements(tmp)
    print(', '.join(get_elements(tmp)))

该方案无需重复回退到文件头，完全符合你提到的「避免重复seek」的需求。

方案2：每次调用重置文件指针（简单实现、性能稍差） #

如果场景对性能要求不高，无需修改原goto_index函数，每次调用前将文件指针重置到开头即可：

def get_elements(f) :
    elements = []
    for idx_str in ('1','2','3') : 
        f.seek(0) # 每次调用前回到文件头重新计数
        elements.append(goto_index(idx_str,f))
    return elements

两种方案修改后运行都可以得到预期输出：hello world, good morning, 69 420

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