Fix #2111: support pickle for built-in dataclasses (#2114)

* 2111: support pickle for built-in dataclasses * 2111: add changes * 2111: simplify test * return original name + handle similar names * add additional check * fix a misspell * remove useless f-string * cleanup test Co-authored-by: Samuel Colvin <[email protected]>
pydantic · Nov 30, 2020 · b1bb6e0 · b1bb6e0
1 parent 9ae40a2
commit b1bb6e0
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 2 deletions.
diff --git a/changes/2111-aimestereo.md b/changes/2111-aimestereo.md
@@ -0,0 +1 @@
+Allow pickling of `pydantic.dataclasses.dataclass` dynamically created from a built-in `dataclasses.dataclass`.
diff --git a/pydantic/dataclasses.py b/pydantic/dataclasses.py
@@ -119,10 +119,23 @@ def _pydantic_post_init(self: 'Dataclass', *initvars: Any) -> None:
     #   __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
@@ -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 @@ class File:
                 '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 @@ class Config:
     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'