pythonpytestpickle

Pickle works but fails pytest


I've created a module persist containing a function obj_pickle which I'm using to pickle various objects in my project. It work fine but it's failing pytest, returning;

>           pickle.dump(obj, file_handle, protocol=protocol)
E           AttributeError: Can't pickle local object 'test_object.<locals>.TestObject'

Running:

I see that the are similar looking issues tied to multiprocessing but I don't think pytest is exposing that issue. Code below, and many thanks.

# lib.persist.py
from pathlib import Path
import pickle


def obj_pickle(obj: object, dir:Path, protocol: int = pickle.HIGHEST_PROTOCOL) -> None:
    """
    Pickle an object to a byte file.
    """
    if not dir.exists():
        dir.mkdir(parents=True, exist_ok=True)
    path = Path(dir, obj.instance_name + '.pkl')
    with open(path, "wb") as file_handle:
        pickle.dump(obj, file_handle, protocol=protocol)
    print(f"{obj.__class__.__name__} object {obj.instance_name} saved to {path}")
# tests.test_persist.py
from pathlib import Path

import pytest

from lib.persist import obj_pickle

TEST_DIR = Path("test_dir")


@pytest.fixture
def test_object():
    class TestObject():
        def __init__(self, instance_file_name):
            self.instance_file_name = instance_file_name
            self.data = "This is a test object."

    test_object = TestObject("test_object_file")
    
    return test_object


def test_obj_pickle(test_object):
    obj_pickle(test_object, Path(TEST_DIR))
    path = Path(TEST_DIR, "test_object_file" + ".pkl")
    assert path.exists()

Solution

  • This doesn't really have anything to do with pytest.

    When you load an object from a pickle, the pickle machinery needs to know what the class of that object should be. But if you define a class inside a function, then every execution of that function generates a whole new class. pickle has no way to tell which version of the class to use, if any versions even exist in the Python session in which you're loading the pickle.

    Instead of trying to make it work, and creating a pile of weird inconsistencies and edge cases, pickle just doesn't support pickling instances of classes like that.

    You need to move your class definition out of the function you have it in:

    class TestObject:
        def __init__(self, instance_file_name):
            self.instance_file_name = instance_file_name
            self.data = "This is a test object."
    
    @pytest.fixture
    def test_object():
        return TestObject("test_object_file")