Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

👌 IMPROVE: Add dataclass serialisation to context #5833

Merged
merged 3 commits into from
Dec 14, 2022
Merged

👌 IMPROVE: Add dataclass serialisation to context #5833

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
93a91e2
👌 IMPROVE: Add dataclass serialisation to context
chrisjsewell Dec 13, 2022
a7e4004
Merge branch 'main' into serialize-dataclass
chrisjsewell Dec 13, 2022
f236f54
Update tests/orm/utils/test_serialize.py
chrisjsewell Dec 13, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions aiida/orm/utils/serialize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,10 @@
"""
from __future__ import annotations

from dataclasses import asdict, is_dataclass
from enum import Enum
from functools import partial
import inspect
from typing import Any, Protocol, Type, overload

from plumpy import Bundle, get_object_loader # type: ignore[attr-defined]
Expand All @@ -28,6 +30,7 @@
from aiida.common import AttributeDict

_ENUM_TAG = '!enum'
_DATACLASS_TAG = '!dataclass'
_NODE_TAG = '!aiida_node'
_GROUP_TAG = '!aiida_group'
_COMPUTER_TAG = '!aiida_computer'
Expand All @@ -51,6 +54,25 @@ def enum_constructor(loader: yaml.Loader, serialized: yaml.Node) -> Enum:
return enum


def represent_dataclass(dumper: yaml.Dumper, obj: Any) -> yaml.MappingNode:
"""Represent an arbitrary dataclass in yaml."""
loader = get_object_loader()
data = {
'__type__': loader.identify_object(obj.__class__),
'__fields__': asdict(obj),
}
return dumper.represent_mapping(_DATACLASS_TAG, data)


def dataclass_constructor(loader: yaml.Loader, serialized: yaml.Node) -> Any:
"""Construct a dataclass from the serialized representation."""
deserialized = loader.construct_mapping(serialized, deep=True) # type: ignore[arg-type]
identifier = deserialized['__type__']
cls = get_object_loader().load_object(identifier)
data = deserialized['__fields__']
return cls(**data)


def represent_node(dumper: yaml.Dumper, node: orm.Node) -> yaml.ScalarNode:
"""Represent a node in yaml."""
if not node.is_stored:
Expand Down Expand Up @@ -136,6 +158,8 @@ def represent_data(self, data):
return represent_computer(self, data)
if isinstance(data, orm.Group):
return represent_group(self, data)
if is_dataclass(data) and not inspect.isclass(data):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just out of interest, why is the second clause necessary? What is an example of an object passing is_dataclass(obj) and inspect.isclass(obj)?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

In [15]: import inspect
In [16]: from dataclasses import asdict, dataclass, is_dataclass

In [17]: @dataclass
    ...: class DataClass:
    ...:     """A dataclass for testing."""
    ...:     my_value: int
    ...: 

In [18]: is_dataclass(DataClass)  and inspect.isclass(DataClass)
Out[18]: True

In [19]: asdict(DataClass)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-19-c90f518d2aad> in <module>
----> 1 asdict(DataClass)

~/Documents/GitHub/aiida_core_develop/.tox/py38-pre-commit/lib/python3.8/dataclasses.py in asdict(obj, dict_factory)
   1070     """
   1071     if not _is_dataclass_instance(obj):
-> 1072         raise TypeError("asdict() should be called on dataclass instances")
   1073     return _asdict_inner(obj, dict_factory)
   1074 

TypeError: asdict() should be called on dataclass instances

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for some reason, _is_dataclass_instance is private 🤷

return represent_dataclass(self, data)

return super().represent_data(data)

Expand Down Expand Up @@ -163,6 +187,7 @@ class AiiDALoader(yaml.Loader):
yaml.add_constructor(_GROUP_TAG, group_constructor, Loader=AiiDALoader)
yaml.add_constructor(_COMPUTER_TAG, computer_constructor, Loader=AiiDALoader)
yaml.add_constructor(_ENUM_TAG, enum_constructor, Loader=AiiDALoader)
yaml.add_constructor(_DATACLASS_TAG, dataclass_constructor, Loader=AiiDALoader)


@overload
Expand Down
17 changes: 17 additions & 0 deletions tests/orm/utils/test_serialize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
# For further information please visit http://www.aiida.net #
###########################################################################
"""Tests for the :mod:`aiida.orm.utils.serialize` module."""
from dataclasses import dataclass
import types
import uuid

Expand Down Expand Up @@ -162,3 +163,19 @@ def test_enum():

deserialized = serialize.deserialize_unsafe(serialized)
assert deserialized == enum


@dataclass
class DataClass:
"""A dataclass for testing."""
my_value: int


def test_dataclass():
"""Test serialization and deserialization of a ``dataclass``."""
obj = DataClass(1)
serialized = serialize.serialize(obj)
assert isinstance(serialized, str)

deserialized = serialize.deserialize_unsafe(serialized)
assert deserialized == obj