pyfunction handles already serialized results (#23)

Author: superstar54

src/aiida_pythonjob/calculations/pyfunction.py

@@ -202,8 +202,10 @@ def run(self) -> ExitCode | None:
                 if exit_code.status != 0:
                     return exit_code
             if len(top_level_output_list) == 1:
-                # If output name in results, use it
-                if top_level_output_list[0]["name"] in results:
+                # User returned a single (nested) dict with AiiDA data nodes as values
+                if self.already_serialized(results):
+                    top_level_output_list = [{"name": key, "value": value} for key, value in results.items()]
+                elif top_level_output_list[0]["name"] in results:
                     top_level_output_list[0]["value"] = self.serialize_output(
                         results.pop(top_level_output_list[0]["name"]),
                         top_level_output_list[0],
@@ -228,8 +230,14 @@ def run(self) -> ExitCode | None:
                 if len(results) > 0:
                     self.logger.warning(f"Found extra results that are not included in the output: {results.keys()}")
         elif len(top_level_output_list) == 1:
-            # Single top-level output, single result
-            top_level_output_list[0]["value"] = self.serialize_output(results, top_level_output_list[0])
+            # Single top-level output
+            # There are two cases:
+            # 1. The output as a whole will be serialized as the single output
+            # 2. The output is a mapping with already AiiDA data nodes as values, no need to serialize
+            if self.already_serialized(results):
+                top_level_output_list[0]["value"] = results
+            else:
+                top_level_output_list[0]["value"] = self.serialize_output(results, top_level_output_list[0])
         else:
             return self.exit_codes.ERROR_RESULT_OUTPUT_MISMATCH
         # Store the outputs
@@ -238,6 +246,20 @@ def run(self) -> ExitCode | None:
 
         return ExitCode()
 
+    def already_serialized(self, results):
+        """Check if the results are already serialized."""
+        import collections
+
+        if isinstance(results, Data):
+            return True
+        elif isinstance(results, collections.abc.Mapping):
+            for value in results.values():
+                if not self.already_serialized(value):
+                    return False
+            return True
+        else:
+            return False
+
     def find_output(self, name):
         """Find the output spec with the given name."""
         for output in self.output_list:

tests/test_pyfunction.py

@@ -1,3 +1,4 @@
+from aiida import orm
 from aiida.engine import run_get_node
 from aiida_pythonjob import pyfunction
 
@@ -118,3 +119,15 @@ def myfunc(x, y):
     assert result["add_multiply"]["add"]["order1"].value == 3
     assert result["add_multiply"]["add"]["order2"].value == 5
     assert result["add_multiply"]["multiply"].value == 2
+
+
+def test_aiida_node_as_inputs_outputs():
+    """Test function with AiiDA nodes as inputs and outputs."""
+
+    @pyfunction()
+    def add(x, y):
+        return {"sum": orm.Int(x + y), "diff": orm.Int(x - y)}
+
+    result, node = run_get_node(add, x=orm.Int(1), y=orm.Int(2))
+    assert set(result.keys()) == {"sum", "diff"}
+    assert result["sum"].value == 3