1  # Simple example presenting how persistent ID can be used to pickle
       2  # external objects by reference.
       3  
       4  import pickle
       5  import sqlite3
       6  from collections import namedtuple
       7  
       8  # Simple class representing a record in our database.
       9  MemoRecord = namedtuple("MemoRecord", "key, task")
      10  
      11  class ESC[4;38;5;81mDBPickler(ESC[4;38;5;149mpickleESC[4;38;5;149m.ESC[4;38;5;149mPickler):
      12  
      13      def persistent_id(self, obj):
      14          # Instead of pickling MemoRecord as a regular class instance, we emit a
      15          # persistent ID.
      16          if isinstance(obj, MemoRecord):
      17              # Here, our persistent ID is simply a tuple, containing a tag and a
      18              # key, which refers to a specific record in the database.
      19              return ("MemoRecord", obj.key)
      20          else:
      21              # If obj does not have a persistent ID, return None. This means obj
      22              # needs to be pickled as usual.
      23              return None
      24  
      25  
      26  class ESC[4;38;5;81mDBUnpickler(ESC[4;38;5;149mpickleESC[4;38;5;149m.ESC[4;38;5;149mUnpickler):
      27  
      28      def __init__(self, file, connection):
      29          super().__init__(file)
      30          self.connection = connection
      31  
      32      def persistent_load(self, pid):
      33          # This method is invoked whenever a persistent ID is encountered.
      34          # Here, pid is the tuple returned by DBPickler.
      35          cursor = self.connection.cursor()
      36          type_tag, key_id = pid
      37          if type_tag == "MemoRecord":
      38              # Fetch the referenced record from the database and return it.
      39              cursor.execute("SELECT * FROM memos WHERE key=?", (str(key_id),))
      40              key, task = cursor.fetchone()
      41              return MemoRecord(key, task)
      42          else:
      43              # Always raises an error if you cannot return the correct object.
      44              # Otherwise, the unpickler will think None is the object referenced
      45              # by the persistent ID.
      46              raise pickle.UnpicklingError("unsupported persistent object")
      47  
      48  
      49  def main():
      50      import io
      51      import pprint
      52  
      53      # Initialize and populate our database.
      54      conn = sqlite3.connect(":memory:")
      55      cursor = conn.cursor()
      56      cursor.execute("CREATE TABLE memos(key INTEGER PRIMARY KEY, task TEXT)")
      57      tasks = (
      58          'give food to fish',
      59          'prepare group meeting',
      60          'fight with a zebra',
      61          )
      62      for task in tasks:
      63          cursor.execute("INSERT INTO memos VALUES(NULL, ?)", (task,))
      64  
      65      # Fetch the records to be pickled.
      66      cursor.execute("SELECT * FROM memos")
      67      memos = [MemoRecord(key, task) for key, task in cursor]
      68      # Save the records using our custom DBPickler.
      69      file = io.BytesIO()
      70      DBPickler(file).dump(memos)
      71  
      72      print("Pickled records:")
      73      pprint.pprint(memos)
      74  
      75      # Update a record, just for good measure.
      76      cursor.execute("UPDATE memos SET task='learn italian' WHERE key=1")
      77  
      78      # Load the records from the pickle data stream.
      79      file.seek(0)
      80      memos = DBUnpickler(file, conn).load()
      81  
      82      print("Unpickled records:")
      83      pprint.pprint(memos)
      84  
      85  
      86  if __name__ == '__main__':
      87      main()