pickleTrial.py

import pickle
import sqlite3
from collections import namedtuple

#Simple class representing a record in our database
MemoRecord=namedtuple("MemoRecord", "key, task")

class DBPickler(pickle.Pickler):
    def persistent_id(self,obj):
        #Instead of pickling MemoRecord as a regular class instance, we emit a
        # persistent ID
        if isinstance(obj, MemoRecord):
            #Here, our persistent ID is simply a tuple, containing a tag and a
            #key, which refers to a specific record in the database
            return ("MemoRecord",obj.key)
        else:
            #If obj does not have a persistent ID, return None. This means obj
            #needs to be pickled as usual
            return None


class DBUnpickler(pickle.Unpickler):
    def __init__(self,file,connection):
        super().__init__(file)
        self.connection=connection

    def persistent_load(self,pid):
        #This method is invoked whenever a persistent ID is encountered
        #Here, pid is the tuple returned by DBPickler
        cursor=self.connection.cursor()
        type_tag, key_id=pid
        if type_tag=="MemoRecord":
            #Fetch the regerenced record from the database and return it
            cursor.execute("SELECT * FROM memos WHERE key=?",(str(key_id),))
            key, task=cursor.fetchone()
            return MemoRecord(key, task)
        else:
            # Always raises an error if you cannot return the correct object
            # Otherwise, the unpickler will think None is the object referenced
            # by the persistent ID
            raise pickle.UnpicklingError("unsupported persistent object")



def main():
    import io
    import pprint

    # Initialize and populate our database
    conn=sqlite3.connect(":memory:")
    cursor=conn.cursor()
    cursor.execute("CREATE TABLE memos(key INTEGER PRIMARY KEY, task TEXT)")
    tasks=(
        'give food to fish',
        'prepare group meeting',
        'fight with a zebra',
        )
    for task in tasks:
        cursor.execute("INSERT INTO memos VALUES(NULL,?)",(task,))

    #Fetch the records to be pickled
    cursor.execute("SELECT * FROM memos")
    memos=[MemoRecord(key,task) for key, task in cursor]
    #Save the records using our custom DBPickler
    file =io.BytesIO()
    DBPickler(file).dump(memos)

    print("Pickled records:")
    pprint.pprint(memos)

    #Update a record, just for good measure
    cursor.execute("UPDATE memos SET task='learn italian' WHERE key=1")

    #Load the records from the pickle data stream
    file.seek(0)
    memos=DBUnpickler(file, conn).load()

    print("Unpickled records:")
    pprint.pprint(memos)


if __name__=='__main__':
    main()