An iterator version of bson.decode_all

bson_decode_iter.py

#coding: utf8
'''
bson_decode_iter.py
~~~~~~~~~~~~~~~~~~~~~~

An iterator version of bson.decode_all

:copyright: (c) 2013 by Yue Du.
:license: The BSD 2-Clause License
'''

import os
import struct
from bson.errors import InvalidBSON
from bson import OLD_UUID_SUBTYPE, ZERO, _elements_to_dict

def decode_iter(fp, as_class=dict, tz_aware=True, uuid_subtype=OLD_UUID_SUBTYPE):
    fp.seek(0, os.SEEK_END)
    length = fp.tell()
    end = length - 1
    position = 0
    while position < end:
        fp.seek(position, os.SEEK_SET)
        obj_size = struct.unpack('<i', fp.read(4))[0]
        if length - position < obj_size:
            raise InvalidBSON("objsize too large")
        fp.seek(position + obj_size - 1, os.SEEK_SET)
        if fp.read(1) != ZERO:
            raise InvalidBSON("bad eoo")
        fp.seek(position + 4, os.SEEK_SET)
        elements = fp.read(obj_size - 5)
        position += obj_size
        yield _elements_to_dict(elements, as_class, tz_aware, uuid_subtype)

if __name__ == '__main__':
    import sys
    with open(sys.argv[1], 'rb') as fp:
        for d in decode_iter(fp):
            print(d)