aiidateam / aiida-core

The official repository for the AiiDA code
https://aiida-core.readthedocs.io
Other
434 stars 188 forks source link

Add support for `JsonableData` to `to_aiida_type` #5721

Closed sphuber closed 1 year ago

sphuber commented 2 years ago

The JsonableData plugin was added to make it easier for users to start wrapping existing Python code with AiiDA. The calcfunction makes it easy to wrap existing Python functions, but they require the arguments to be passed as Data nodes. Various Python base types are now automatically converted through the to_aiida_type (such as int, str, bool etc.) but any other objects need to be manually wrapped.

The JsonableData made this easier for objects that implement an as_dict method, such as a lot of the pymatgen data classes, e.g., Molecule and Structure. It would be great if the to_aiida_type would also automatically serialize these, such that another step of explicit conversion is removed.

sphuber commented 1 year ago

Maybe this isn't a good idea to add after all. I made an implementation, but think it may not be worth to add this. An example usage would look like:

from aiida.engine import calcfunction
from aiida.orm import StructureData
from pymatgen.core import Molecule

molecule = Molecule(['H'], [[0, 0, 0]])

@calcfunction
def convert_pymatgen_to_structure_data(molecule):
    return StructureData(pymatgen=molecule.obj)

structure = convert_pymatgen_to_structure_data(molecule)

Although useful, the solution without this feature would not be all that more cumbersome.

from aiida.engine import calcfunction
from aiida.orm import JsonableData, StructureData
from pymatgen.core import Molecule

molecule = Molecule(['H'], [[0, 0, 0]])

@calcfunction
def convert_pymatgen_to_structure_data(molecule):
    return StructureData(pymatgen=molecule.obj)

structure = convert_pymatgen_to_structure_data(JsonableData(molecule))

The user merely has to wrap the molecule into the JsonableData for it to work. Making the user explicitly do this also makes it a lot easier to spot why inside the function, one has to call molecule.obj to retrieve the originally wrapped object. With the automatic serialization, this may cause confusion. More so than with the automatic serialization of Python base types, since for most base type (int, str, bool, etc) most of the important operations are overloaded and so there is no explicit dereferencing necessary (at least not always).

sphuber commented 1 year ago

Just for posterity, if we still decide to add this, the following patch needs to be applied.

From 2b7341a7f28a6e1af52bef20b6bac0e7b00ef5db Mon Sep 17 00:00:00 2001
From: Sebastiaan Huber <mail@sphuber.net>
Date: Wed, 26 Oct 2022 18:13:31 +0200
Subject: [PATCH] Add support for `JsonableData` to `to_aiida_type`

---
 aiida/orm/nodes/data/base.py               |  7 +++++++
 aiida/orm/nodes/data/jsonable.py           | 12 +++++++++++-
 tests/orm/nodes/data/test_to_aiida_type.py |  2 ++
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/aiida/orm/nodes/data/base.py b/aiida/orm/nodes/data/base.py
index f95cacaa2..7829a0a39 100644
--- a/aiida/orm/nodes/data/base.py
+++ b/aiida/orm/nodes/data/base.py
@@ -18,6 +18,13 @@ __all__ = ('BaseType', 'to_aiida_type')
 @singledispatch
 def to_aiida_type(value):
     """Turns basic Python types (str, int, float, bool) into the corresponding AiiDA types."""
+    from .jsonable import JsonableData
+
+    try:
+        return JsonableData(value)
+    except TypeError:
+        pass
+
     raise TypeError(f'Cannot convert value of type {type(value)} to AiiDA type.')

diff --git a/aiida/orm/nodes/data/jsonable.py b/aiida/orm/nodes/data/jsonable.py
index d16670a4e..e4214fa5a 100644
--- a/aiida/orm/nodes/data/jsonable.py
+++ b/aiida/orm/nodes/data/jsonable.py
@@ -25,7 +25,7 @@ class JsonableData(Data):
     store an instance as a ``JsonableData`` simply pass an instance as an argument to the constructor as follows::

         from pymatgen.core import Molecule
-        molecule = Molecule(['H']. [0, 0, 0])
+        molecule = Molecule(['H'], [[0, 0, 0]])
         node = JsonableData(molecule)
         node.store()

@@ -83,6 +83,16 @@ class JsonableData(Data):

         self.base.attributes.set_many(serialized)

+    def __eq__(self, other):
+        """Return whether this instance is equal to ``other``."""
+        if isinstance(other, self.__class__):
+            return self.base.attributes.all == other.base.attributes.all
+
+        try:
+            return self.obj == other
+        except Exception:  # pylint: disable=broad-except
+            return False
+
     @classmethod
     def _deserialize_float_constants(cls, data: typing.Any):
         """Deserialize the contents of a dictionary ``data`` deserializing infinity and NaN string constants.
diff --git a/tests/orm/nodes/data/test_to_aiida_type.py b/tests/orm/nodes/data/test_to_aiida_type.py
index aa76e79e9..ca02b874e 100644
--- a/tests/orm/nodes/data/test_to_aiida_type.py
+++ b/tests/orm/nodes/data/test_to_aiida_type.py
@@ -8,6 +8,7 @@
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
 """Test the :meth:`aiida.orm.data.base.to_aiida_type` serializer."""
+from pymatgen.core import Molecule
 import pytest

 from aiida import orm
@@ -24,6 +25,7 @@ from aiida.common.links import LinkType
         (orm.List, [0, 1, 2]),
         (orm.Str, 'test-string'),
         (orm.EnumData, LinkType.RETURN),
+        (orm.JsonableData, Molecule(['H'], [[0, 0, 0]]))
     )
 )
 # yapf: enable
-- 
2.25.1