python读取pk文件

import pickle
import os
#获取文件目录
def getFilename():
    dir = os.getcwd()
    listfile = os.listdir(dir)
    test_filename,train_filename = '',''
    for i in range(len(listfile)):
        if listfile[i].find('test.pk') > 0:  # 没找到子串返回-1，找到返回索引位置
            test_filename = dir + '\\' + listfile[i]
        if listfile[i].find('train.pk')>0:
            train_filename=dir+'\\'+listfile[i]
    return test_filename,train_filename

#将读到的pk文件的数据写入新文件
def write_TestdataOrTraindata(filename):
    with open(filename, 'rb') as f:
        data = pickle.load(f)
    f=open(filename[filename.find('_t')+1:filename.find('.')]+'.txt','w',encoding='utf-8')
    for i in range(len(data)):
        print(data[i])
        f.write(str(data[i]))
        f.write('\n')

if __name__=='__main__':
    test_filename,train_filename=getFilename()
    write_TestdataOrTraindata(test_filename)