From 2fa7a530511a94ead83d79669efed71706a0a472 Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Tue, 4 Jun 2024 11:33:11 +0200 Subject: [PATCH] Engine: Set the `to_aiida_type` as default inport port serializer (#6439) The `to_aiida_type` serializer automatically converts a number of Python base types to the corresponding `Data` node when passed as an input to a process. Process functions already automatically set this serializer as a default, which simplifies the life of users and developers. Here, the `to_aiida_type` is now set as the default serializer for all input ports. The only exception is if the port is a metadata input port, in which case the data is stored directly in the attributes of the process node and so should not be converted to a node. It is also skipped if the port itself already declares a serializer. --- src/aiida/engine/processes/ports.py | 18 +++++++++++++++++- tests/engine/test_ports.py | 25 ++++++++++++++++++++++++- tests/engine/test_process.py | 26 ++++++++++++++++++++++++++ 3 files changed, 67 insertions(+), 2 deletions(-) diff --git a/src/aiida/engine/processes/ports.py b/src/aiida/engine/processes/ports.py index 538e6ab195..50f9b0d6e1 100644 --- a/src/aiida/engine/processes/ports.py +++ b/src/aiida/engine/processes/ports.py @@ -8,6 +8,8 @@ ########################################################################### """AiiDA specific implementation of plumpy Ports and PortNamespaces for the ProcessSpec.""" +from __future__ import annotations + import re import warnings from collections.abc import Mapping @@ -17,7 +19,7 @@ from plumpy.ports import breadcrumbs_to_port from aiida.common.links import validate_link_label -from aiida.orm import Data, Node +from aiida.orm import Data, Node, to_aiida_type __all__ = ( 'PortNamespace', @@ -115,6 +117,11 @@ def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self._serializer: Callable[[Any], 'Data'] = serializer + @property + def serializer(self) -> Callable[[Any], 'Data'] | None: + """Return the serializer.""" + return self._serializer + def serialize(self, value: Any) -> 'Data': """Serialize the given value, unless it is ``None``, already a Data type, or no serializer function is defined. @@ -209,6 +216,15 @@ def __setitem__(self, key: str, port: ports.Port) -> None: if hasattr(port, 'non_db_explicitly_set') and not port.non_db_explicitly_set: # type: ignore[attr-defined] port.non_db = self.non_db # type: ignore[attr-defined] + # If the port is not metadata (signified by ``is_metadata`` and ``non_db`` being ``False`` if defined) and it + # does not already define a serializer, set the default serializer to ``to_aiida_type``. + if ( + ((hasattr(port, 'is_metadata') and not port.is_metadata) and (hasattr(port, 'non_db') and not port.non_db)) + and hasattr(port, 'serializer') + and port.serializer is None + ): + port._serializer = to_aiida_type + super().__setitem__(key, port) @staticmethod diff --git a/tests/engine/test_ports.py b/tests/engine/test_ports.py index 83134d6d7b..147717f8a4 100644 --- a/tests/engine/test_ports.py +++ b/tests/engine/test_ports.py @@ -10,7 +10,7 @@ import pytest from aiida.engine.processes.ports import InputPort, PortNamespace -from aiida.orm import Dict, Int +from aiida.orm import Dict, Int, to_aiida_type class TestInputPort: @@ -119,3 +119,26 @@ def test_lambda_default(self): inputs = port_namespace.pre_process({'port': Int(3)}) assert isinstance(inputs['port'], Int) assert inputs['port'].value == 3 + + @pytest.mark.parametrize('is_metadata', (True, False)) + def test_serializer(self, is_metadata): + """Test the automatic setting of the serializer for non-metadata ports. + + For non-metadata portnamespaces, the serializer should be automatically set to the + :meth:`aiida.orm.nodes.data.base.to_aiida_type` serializer, unless the port already defines a serializer itself. + """ + namespace = PortNamespace('namespace', is_metadata=is_metadata) + port = InputPort('port') + namespace['port'] = port + + if is_metadata: + assert port.serializer is None + else: + assert port.serializer is to_aiida_type + + def custom_serializer(*args, **kwargs): + pass + + other_port = InputPort('other_port', serializer=custom_serializer) + namespace['other_port'] = port + assert other_port.serializer is custom_serializer diff --git a/tests/engine/test_process.py b/tests/engine/test_process.py index 07ac4c4c8b..3035b432b7 100644 --- a/tests/engine/test_process.py +++ b/tests/engine/test_process.py @@ -17,6 +17,7 @@ from aiida.engine import ExitCode, ExitCodesNamespace, Process, run, run_get_node, run_get_pk from aiida.engine.processes.ports import PortNamespace from aiida.manage.caching import disable_caching, enable_caching +from aiida.orm import to_aiida_type from aiida.orm.nodes.caching import NodeCaching from aiida.plugins import CalculationFactory from plumpy.utils import AttributesFrozendict @@ -568,3 +569,28 @@ def test_metadata_disable_cache(runner, entry_points): with enable_caching(): process = CachableProcess(runner=runner, inputs={'metadata': {'disable_cache': True}}) assert not process.node.base.caching.is_created_from_cache + + +def custom_serializer(): + pass + + +class AutoSerializeProcess(Process): + """Check the automatic assignment of ``to_aiida_type`` serializer.""" + + @classmethod + def define(cls, spec): + super().define(spec) + spec.input('non_metadata_input') + spec.input('metadata_input', is_metadata=True) + spec.input('custom_input', serializer=custom_serializer) + + +def test_auto_default_serializer(): + """Test that all inputs ports automatically have ``to_aiida_type`` set as the serializer. + + Exceptions are if the port is a metadata port or it defines an explicit serializer + """ + assert AutoSerializeProcess.spec().inputs['non_metadata_input'].serializer is to_aiida_type + assert AutoSerializeProcess.spec().inputs['metadata_input'].serializer is None + assert AutoSerializeProcess.spec().inputs['custom_input'].serializer is custom_serializer