MongoDB lazy saving

autosave.py

class AutoSave(object):
 
    def __init__(self, d, collection):
        self.d = d
        self.collection = collection
 
    def __enter__(self):
        return self.d
 
    def __exit__(self, type, value, traceback):
        print type
        print value
        print traceback
        if self.d and not type:
            self.collection.save(self.d)

class make_autosave(object):
    def __init__(self, collection):
        self.collection = collection

    def find_one(self, d=None):
        obj = self.collection.find_one(d)
        return AutoSave(obj, self.collection)

    def insert(self, d=None):
        _id = self.collection.insert(d if d else {})
        return self.find_one({'_id': _id})

test.py

import autosave
from pymongo import MongoClient

# Setup
TEST_DB = 'test_mongoautosave'

conn = MongoClient()
conn.drop_database(TEST_DB)
db = conn[TEST_DB]
col1 = db.col1
col2 = db.col2

autosave_col1 = autosave.make_autosave(col1)
autosave_col2 = autosave.make_autosave(col2)

def test_1():
    with autosave.AutoSave({}, col1) as d:
        d['row1'] = 'row1'
    assert col1.find_one({'row1': 'row1'})

def test_2():
    with autosave_col1.find_one({'row1':'row1'}) as d:
        d['row2'] = 'row2'

def test_3():
    with autosave_col2.insert() as d:
        assert '_id' in d