Сохранение внешних объектов в pickle

Для удобства сохранения объекта, модуль pickle поддерживает понятие ссылки на объект вне потока обработанных данных. На такие объекты ссылается постоянный идентификатор, который должен быть либо строкой буквенно-цифровых символов для протокола 0, либо просто произвольным объектом для любого более нового протокола.

Разрешение таких постоянных идентификаторов не определяется модулем pickle. Он будет делегировать это разрешение пользовательским методам во время сохранения - persistent_id() и во время извлечения - persistent_load() соответственно.

Для выбора объектов, имеющих внешний постоянный идентификатор, pickle.Pickler() должен иметь собственный метод Pickler.persistent_id(obj), который принимает объект obj в качестве аргумента и возвращает либо None, либо постоянный идентификатор для этого объекта. Когда возвращается None, то pickle просто выбирает объект как обычно. Когда возвращается постоянная строка идентификатора, pickle выберет этот объект вместе с маркером, чтобы распознать его как постоянный идентификатор.

При распаковке, чтобы открепить внешние объекты, у разборщика pickle.Unpickler() должен быть собственный метод Unpickler.persistent_load(), который принимает постоянный объект ID и возвращает указанный объект.

Простой пример, показывающий, как постоянный идентификатор можно использовать для выбора внешних объектов по ссылке.

import pickle
import sqlite3
from collections import namedtuple

# Простой класс, представляющий запись в базе данных.
MemoRecord = namedtuple("MemoRecord", "key, task")

class DBPickler(pickle.Pickler):

    def persistent_id(self, obj):
        # Генерируем постоянный идентификатор, вместо того, 
        # чтобы выбирать MemoRecord как обычный экземпляр класса
        if isinstance(obj, MemoRecord):
            # Здесь постоянный идентификатор - это просто кортеж, 
            # содержащий тег и ключ, который ссылается на конкретную
            # запись в базе данных.
            return ("MemoRecord", obj.key)
        else:
            # Если у 'obj' нет постоянного идентификатора, то вернем 'None'. 
            # Это означает, что 'obj' нужно сохранять как обычно.
            return None


class DBUnpickler(pickle.Unpickler):

    def __init__(self, file, connection):
        super().__init__(file)
        self.connection = connection

    def persistent_load(self, pid):
        # Этот метод вызывается всякий раз, когда встречается постоянный
        # идентификатор. Здесь pid - это кортеж, возвращаемый 'DBPickler'.
        cursor = self.connection.cursor()
        type_tag, key_id = pid
        if type_tag == "MemoRecord":
            # Получить указанную запись из базы данных и возвращаем ее.
            cursor.execute("SELECT * FROM memos WHERE key=?", (str(key_id),))
            key, task = cursor.fetchone()
            return MemoRecord(key, task)
        else:
            # Всегда выдает ошибку, если нельзя вернуть правильный объект. 
            # В противном случае 'Unpickler' будет думать, что 'None' - это объект, 
            # на который ссылается постоянный идентификатор.
            raise pickle.UnpicklingError("unsupported persistent object")


def main():
    import io
    import pprint

    # Инициализируем и заполним базу данных.
    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,))

    # Получить записи для сохранения.
    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)

    # Сохраним записи, используя собственный 'DBPickler'.
    cursor.execute("UPDATE memos SET task='learn italian' WHERE key=1")

    # Загрузим записи из потока данных 'pickle'.
    file.seek(0)
    memos = DBUnpickler(file, conn).load()

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


if __name__ == '__main__':
    main()