Python 删减pkl内容

import os
import pickle
import numpy as np

def read_pickle(work_path):
 	'''
    todo:   read pkl to list
    '''
    data_list = []
    with open(work_path, "rb") as f:
            while True:
                try:
                    data = pickle.load(f)
                    if data['image_path'] == 'xxxx':
                        continue
                    data_list.append(data)
                except EOFError:
                    break
    return data_list

def write_pkl(save_path, image_list):
    '''
    todo:   write list to pkl
    '''
    train_pkl_path = os.path.join(save_path, 'train.pkl') # can change file name str
    with open(train_pkl_path, "wb") as train_pkl:
        for i in image_list:
            pickle.dump(i, train_pkl)
       

def main():
    pkl_path = '/home/xxx/train.pkl'
    data_list = read_pickle(pkl_path)
    write_pkl(os.path.dirname(pkl_path), data_list)

if __name__=="__main__":
    main()