python之Unpickling 中流(python)

我正在编写脚本来通过反复取消对象直到 EOF 来处理(非常大的)文件。我想对文件进行分区并让单独的进程(在云中)解开并处理单独的部分。

但是我的分区器并不智能，它不知道文件中 pickle 对象之间的边界(因为这些边界取决于被 pickle 的对象类型等)。

有没有办法扫描文件中的“开始 pickle 对象”哨兵？天真的方法是尝试在连续的字节偏移处进行 unpickling，直到对象被成功 pickle 为止，但这会产生意想不到的错误。似乎对于某些输入组合，unpickler 不同步并且不会为文件的其余部分返回任何内容(请参见下面的代码)。

import cPickle 
import os 
 
def stream_unpickle(file_obj): 
    while True: 
        start_pos = file_obj.tell() 
        try: 
            yield cPickle.load(file_obj) 
        except (EOFError, KeyboardInterrupt): 
            break 
        except (cPickle.UnpicklingError, ValueError, KeyError, TypeError, ImportError): 
            file_obj.seek(start_pos+1, os.SEEK_SET) 
 
if __name__ == '__main__': 
    import random 
    from StringIO import StringIO 
 
    # create some data 
    sio = StringIO() 
    [cPickle.dump(random.random(), sio, cPickle.HIGHEST_PROTOCOL) for _ in xrange(1000)] 
    sio.flush() 
 
    # read from subsequent offsets and find discontinuous jumps in object count 
    size = sio.tell() 
    last_count = None 
    for step in xrange(size): 
        sio.seek(step, os.SEEK_SET) 
        count = sum(1 for _ in stream_unpickle(file_obj)) 
        if last_count is None or count == last_count - 1: 
            last_count = count 
        elif count != last_count: 
            # if successful, these should never print (but they do...) 
            print '%d elements read from byte %d' % (count, step) 
            print '(%d elements read from byte %d)' % (last_count, step-1) 
            last_count = count 

请您参考如下方法:

pickletools 模块有一个显示操作码的 dis 函数。它表明有一个您可能正在扫描的 STOP 操作码:

>>> import pickle, pickletools, StringIO 
>>> s = StringIO.StringIO() 
>>> pickle.dump('abc', s) 
>>> p = s.getvalue() 
>>> pickletools.dis(p) 
    0: S    STRING     'abc' 
    7: p    PUT        0 
   10: .    STOP 
highest protocol among opcodes = 0 

请注意，使用 STOP 操作码有点棘手，因为代码的长度可变，但它可以作为有关截止位置的有用提示。

如果您控制另一端的 pickle 步骤，则可以通过添加您自己明确的替代分隔符来改善这种情况:

>>> sep = '\xDE\xAD\xBE\xEF' 
>>> s = StringIO.StringIO() 
>>> pickle.dump('abc', s) 
>>> s.write(sep) 
>>> pickle.dump([10, 20], s) 
>>> s.write(sep) 
>>> pickle.dump('def', s) 
>>> s.write(sep) 
>>> pickle.dump([30, 40], s) 
>>> p = s.getvalue() 

在拆包之前，使用已知的分隔符分成单独的 pickle :

>>> for pick in p.split(sep): 
        print pickle.loads(pick) 
 
abc 
[10, 20] 
def 
[30, 40]