Pydantic: How to pickle a generic pydantic class?

Created on 28 Jun 2020  路  5Comments  路  Source: samuelcolvin/pydantic

Hi,

I am trying to pickle a pydantic generic model as illustrated below, using Python 3.7.6 and Pydantic 1.5.1:

import pickle
from typing import TypeVar, Generic
from pydantic import BaseModel
from pydantic.generics import GenericModel

T = TypeVar('T')


class MyClass(BaseModel):
    pass


class GenericClass(GenericModel, Generic[T]):
    value: T


obj = MyClass()
instance = GenericClass[MyClass](value=obj)

pickle.dumps(instance)

This piece of code generates this error:

_pickle.PicklingError: Can't pickle <class '__main__.GenericClass[MyClass]'>: 
attribute lookup GenericClass[MyClass] on __main__ failed

It seems to be an error related specifically to pydantic generic models, as I am able to pickle non generic pydantic objects and generic objects that are not pydantic.

Am I doing something wrong or is this not supported yet?
Thanks for your help!

bug

All 5 comments

Looks like it's not currently supported.

Should be possible to fix, but it's not a priority for me.

Hi Samuel,

Thanks for your answer and for your great work.
Looking forward to the fix when time allows,

Jo毛l.

I was able to solve this problem by setting the attribute with new class to the module. Will make a PR soon.

As a hotfix, you can try this:

```py
import pickle
import sys
from typing import TypeVar, Generic, Type, Union, Any, Tuple

from pydantic import BaseModel
from pydantic.generics import GenericModel, GenericModelT

T = TypeVar('T')

class MyClass(BaseModel): ...

class GenericClass(GenericModel, Generic[T]):
value: T

def __class_getitem__(cls: Type[GenericModelT], params: Union[Type[Any], Tuple[Type[Any], ...]]) -> Type[Any]:
    created_class = super().__class_getitem__(params)
    setattr(sys.modules[created_class.__module__], created_class.__name__, created_class)
    return created_class

obj = MyClass()
instance = GenericClassMyClass
print(pickle.loads(pickle.dumps(instance)))

> value=MyClass()

Found out that this issue relates not only to generic models, but to any dynamically created model not listed in its module globals.

This would be a great feature to add as the inability to pickle generic classes prohibits using these classes for deep learning configurations as multi-gpu setups like PyTorch Dataparallel require pickling capacity of the models.

Was this page helpful?
0 / 5 - 0 ratings

Related issues

AlbertMukhammadiev picture AlbertMukhammadiev  路  3Comments

iwoloschin picture iwoloschin  路  3Comments

mgresko picture mgresko  路  3Comments

samuelcolvin picture samuelcolvin  路  3Comments

cdeil picture cdeil  路  3Comments