Skip to content

Commit

Permalink
A temporary quick fix for dataclass serialization (#500)
Browse files Browse the repository at this point in the history
This quick fix will be removed when proper dataclass serialization 
support is added to dill. This is just here to allow for better support, 
at least for now. dataclasses pickled with this PR will be unpicklable 
by future versions of dill, but the future versions of dill will be able 
to be automatically use the newer features in dataclasses.py that were 
not available in older versions of Python. That forward compatibility 
features is not present in this PR.
  • Loading branch information
anivegesana committed Jun 8, 2022
1 parent 4462409 commit 1f5814e
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 0 deletions.
31 changes: 31 additions & 0 deletions dill/_dill.py
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,10 @@ def get_file_type(*args, **kwargs):
from collections import OrderedDict

import inspect
import dataclasses

from pickle import GLOBAL


### Shims for different versions of Python and dill
class Sentinel(object):
Expand Down Expand Up @@ -2258,6 +2262,33 @@ def save_capsule(pickler, obj):
else:
_testcapsule = None


#############################
# A quick fix for issue #500
# This should be removed when a better solution is found.

if hasattr(dataclasses, "_HAS_DEFAULT_FACTORY_CLASS"):
@register(dataclasses._HAS_DEFAULT_FACTORY_CLASS)
def save_dataclasses_HAS_DEFAULT_FACTORY_CLASS(pickler, obj):
pickler.write(GLOBAL + b"dataclasses\n_HAS_DEFAULT_FACTORY\n")

if hasattr(dataclasses, "MISSING"):
@register(type(dataclasses.MISSING))
def save_dataclasses_MISSING_TYPE(pickler, obj):
pickler.write(GLOBAL + b"dataclasses\nMISSING\n")

if hasattr(dataclasses, "KW_ONLY"):
@register(type(dataclasses.KW_ONLY))
def save_dataclasses_KW_ONLY_TYPE(pickler, obj):
pickler.write(GLOBAL + b"dataclasses\nKW_ONLY\n")

if hasattr(dataclasses, "_FIELD_BASE"):
@register(dataclasses._FIELD_BASE)
def save_dataclasses_FIELD_BASE(pickler, obj):
pickler.write(GLOBAL + b"dataclasses\n" + obj.name.encode() + b"\n")

#############################

# quick sanity checking
def pickles(obj,exact=False,safe=False,**kwds):
"""
Expand Down
35 changes: 35 additions & 0 deletions tests/test_dataclasses.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
#!/usr/bin/env python
#
# Author: Mike McKerns (mmckerns @caltech and @uqfoundation)
# Author: Anirudh Vegesana (avegesan@cs.stanford.edu)
# Copyright (c) 2022 The Uncertainty Quantification Foundation.
# License: 3-clause BSD. The full license text is available at:
# - https://github.com/uqfoundation/dill/blob/master/LICENSE
"""
test pickling a dataclass
"""

import dill
import dataclasses

def test_dataclasses():
# Issue #500
@dataclasses.dataclass
class A:
x: int
y: str

@dataclasses.dataclass
class B:
a: A

a = A(1, "test")
before = B(a)
save = dill.dumps(before)
after = dill.loads(save)
assert before != after # classes don't match
assert before == B(A(**dataclasses.asdict(after.a)))
assert dataclasses.asdict(before) == dataclasses.asdict(after)

if __name__ == '__main__':
test_dataclasses()

0 comments on commit 1f5814e

Please sign in to comment.