Allow passing user when creating AiiDA data (#61)

Support creating an AiiDA node in a different thread.

Author: superstar54

src/aiida_pythonjob/data/jsonable_data.py

@@ -64,19 +64,20 @@ def _extract_dict(self, obj: typing.Any) -> dict:
             f"Object `{obj}` does not have any of the following dictionary-conversion methods: {self._DICT_METHODS}"
         )
 
-    def _make_jsonable(self, data: typing.Any) -> typing.Any:
+    @classmethod
+    def _make_jsonable(cls, data: typing.Any) -> typing.Any:
         """
         Recursively walk `data`. Convert anything that is not JSON-serializable
         into JSON-friendly structures (e.g. convert NumPy arrays to lists).
         """
         if isinstance(data, dict):
-            return {k: self._make_jsonable(v) for k, v in data.items()}
+            return {k: cls._make_jsonable(v) for k, v in data.items()}
 
         elif isinstance(data, list):
-            return [self._make_jsonable(v) for v in data]
+            return [cls._make_jsonable(v) for v in data]
 
         elif isinstance(data, tuple):
-            return tuple(self._make_jsonable(v) for v in data)
+            return tuple(cls._make_jsonable(v) for v in data)
 
         elif isinstance(data, np.ndarray):
             return data.tolist()

src/aiida_pythonjob/data/serializer.py

@@ -103,6 +103,7 @@ def general_serializer(
     data: Any,
     serializers: dict | None = None,
     store: bool = True,
+    user: orm.User | None = None,
 ) -> orm.Node:
     """
     Attempt to serialize the data to an AiiDA data node based on the preference from `config`:
@@ -123,7 +124,7 @@ def general_serializer(
     if ep_key in serializers:
         try:
             serializer = import_from_path(serializers[ep_key])
-            new_node = serializer(data)
+            new_node = serializer(data, user=user)
             if store:
                 new_node.store()
             return new_node
@@ -132,7 +133,7 @@ def general_serializer(
             raise ValueError(f"Error in serializing {ep_key}: {error_traceback}")
 
     try:
-        node = JsonableData(data)
+        node = JsonableData(data, user=user)
         if store:
             node.store()
         return node

src/aiida_pythonjob/parsers/utils.py

@@ -35,6 +35,7 @@ def parse_outputs(
     exit_codes,
     logger,
     serializers: Optional[Dict[str, str]] = None,
+    user: Optional[orm.User] = None,
 ) -> Tuple[Optional[Dict[str, Any]], Optional[ExitCode]]:
     """Validate & convert *results* according to *output_spec*.
 
@@ -58,7 +59,7 @@ def parse_outputs(
         for i, name in enumerate(names):
             child_spec = fields[name]
             val = results[i]
-            outs[name] = serialize_ports(val, child_spec, serializers=serializers)
+            outs[name] = serialize_ports(val, child_spec, serializers=serializers, user=user)
         return outs, None
 
     # dict
@@ -83,19 +84,19 @@ def parse_outputs(
             ((only_name, only_spec),) = fields.items()
             # if user used the same key as port name, use that value;
             if only_name in results:
-                outs[only_name] = serialize_ports(results.pop(only_name), only_spec, serializers=serializers)
+                outs[only_name] = serialize_ports(results.pop(only_name), only_spec, serializers=serializers, user=user)
                 if results:
                     logger.warning(f"Found extra results that are not included in the output: {list(results.keys())}")
             else:
                 # else treat the entire dict as the value for that single port.
-                outs[only_name] = serialize_ports(results, only_spec, serializers=serializers)
+                outs[only_name] = serialize_ports(results, only_spec, serializers=serializers, user=user)
             return outs, None
 
         # fixed fields
         for name, child_spec in fields.items():
             if name in remaining:
                 value = remaining.pop(name)
-                outs[name] = serialize_ports(value, child_spec, serializers=serializers)
+                outs[name] = serialize_ports(value, child_spec, serializers=serializers, user=user)
             else:
                 # If the field is explicitly required -> invalid output
                 required = getattr(child_spec.meta, "required", None)
@@ -109,6 +110,7 @@ def parse_outputs(
                     remaining,
                     spec or SocketSpec(identifier="node_graph.any"),
                     serializers=serializers,
+                    user=user,
                 )
             )
             return outs, None
@@ -120,7 +122,7 @@ def parse_outputs(
     # single fixed output + non-dict/tuple scalar
     if len(fields) == 1 and not is_dyn:
         ((only_name, only_spec),) = fields.items()
-        return {only_name: serialize_ports(results, only_spec, serializers=serializers)}, None
+        return {only_name: serialize_ports(results, only_spec, serializers=serializers, user=user)}, None
 
     # empty output spec + None result
     if len(fields) == 0 and results is None:

src/aiida_pythonjob/utils.py

@@ -283,6 +283,7 @@ def serialize_ports(
     python_data: Any,
     port_schema: SocketSpec | Dict[str, Any],
     serializers: Optional[Dict[str, str]] = None,
+    user: Optional[User] = None,
 ) -> Any:
     """Serialize raw Python data to AiiDA Data following a SocketSpec schema.
 
@@ -309,21 +310,21 @@ def serialize_ports(
                 elif child_spec.is_namespace():
                     out[key] = serialize_ports(value, child_spec, serializers=serializers)
                 else:
-                    out[key] = general_serializer(value, serializers=serializers, store=False)
+                    out[key] = general_serializer(value, serializers=serializers, store=False, user=user)
             elif spec.meta.dynamic:
                 if spec.item is None:
                     if isinstance(value, dict):
                         item = SocketSpec(identifier="node_graph.namespace", meta=SocketMeta(dynamic=True))
                         out[key] = serialize_ports(value, item, serializers=serializers)
                     else:
-                        out[key] = general_serializer(value, serializers=serializers, store=False)
+                        out[key] = general_serializer(value, serializers=serializers, store=False, user=user)
                 elif spec.item.is_namespace():
                     out[key] = serialize_ports(value, spec.item, serializers=serializers)
                 else:
-                    out[key] = general_serializer(value, serializers=serializers, store=False)
+                    out[key] = general_serializer(value, serializers=serializers, store=False, user=user)
             else:
                 raise ValueError(f"Unexpected key '{key}' for namespace '{name}' (not dynamic).")
         return out
 
     # Leaf
-    return general_serializer(python_data, serializers=serializers, store=False)
+    return general_serializer(python_data, serializers=serializers, store=False, user=user)