diff --git a/changes/2111-aimestereo.md b/changes/2111-aimestereo.md new file mode 100644 index 0000000..d930dfc --- /dev/null +++ b/changes/2111-aimestereo.md @@ -0,0 +1 @@ +Allow pickling of `pydantic.dataclasses.dataclass` dynamically created from a built-in `dataclasses.dataclass`. \ No newline at end of file diff --git a/pydantic/dataclasses.py b/pydantic/dataclasses.py index 0d28e1b..53f1427 100644 --- a/pydantic/dataclasses.py +++ b/pydantic/dataclasses.py @@ -119,10 +119,23 @@ def _process_class( # __post_init__ = _pydantic_post_init # ``` # with the exact same fields as the base dataclass + # and register it on module level to address pickle problem: + # https://github.com/samuelcolvin/pydantic/issues/2111 if is_builtin_dataclass(_cls): + uniq_class_name = f'_Pydantic_{_cls.__name__}_{id(_cls)}' _cls = type( - _cls.__name__, (_cls,), {'__annotations__': _cls.__annotations__, '__post_init__': _pydantic_post_init} + # for pretty output new class will have the name as original + _cls.__name__, + (_cls,), + { + '__annotations__': _cls.__annotations__, + '__post_init__': _pydantic_post_init, + # attrs for pickle to find this class + '__module__': __name__, + '__qualname__': uniq_class_name, + }, ) + globals()[uniq_class_name] = _cls else: _cls.__post_init__ = _pydantic_post_init cls: Type['Dataclass'] = dataclasses.dataclass( # type: ignore diff --git a/tests/test_dataclasses.py b/tests/test_dataclasses.py index 4e71a7b..e087455 100644 --- a/tests/test_dataclasses.py +++ b/tests/test_dataclasses.py @@ -1,4 +1,5 @@ import dataclasses +import pickle from collections.abc import Hashable from datetime import datetime from pathlib import Path @@ -733,7 +734,10 @@ def test_override_builtin_dataclass_nested_schema(): 'type': 'object', } }, - 'properties': {'filename': {'title': 'Filename', 'type': 'string'}, 'meta': {'$ref': '#/definitions/Meta'}}, + 'properties': { + 'filename': {'title': 'Filename', 'type': 'string'}, + 'meta': {'$ref': '#/definitions/Meta'}, + }, 'required': ['filename', 'meta'], 'title': 'File', 'type': 'object', @@ -795,3 +799,37 @@ def test_forward_stdlib_dataclass_params(): e.other = 'bulbi2' with pytest.raises(dataclasses.FrozenInstanceError): e.item.name = 'pika2' + + +def test_pickle_overriden_builtin_dataclass(create_module): + module = create_module( + # language=Python + """\ +import dataclasses +import pydantic + + +@dataclasses.dataclass +class BuiltInDataclassForPickle: + value: int + +class ModelForPickle(pydantic.BaseModel): + # pickle can only work with top level classes as it imports them + + dataclass: BuiltInDataclassForPickle + + class Config: + validate_assignment = True + """ + ) + obj = module.ModelForPickle(dataclass=module.BuiltInDataclassForPickle(value=5)) + + pickled_obj = pickle.dumps(obj) + restored_obj = pickle.loads(pickled_obj) + + assert restored_obj.dataclass.value == 5 + assert restored_obj == obj + + # ensure the restored dataclass is still a pydantic dataclass + with pytest.raises(ValidationError, match='value\n +value is not a valid integer'): + restored_obj.dataclass.value = 'value of a wrong type'