From 1cdaa442e51eb57a2c102c21138c27c5bbc7c2a1 Mon Sep 17 00:00:00 2001
From: Xing Wang <xingwang1991@gmail.com>
Date: Sat, 10 Aug 2024 17:54:50 +0200
Subject: [PATCH] Set the node' label as the workgraph/task's name. (#213)

Set the node' label as the workgraph/task's name when the process node is created. This is very useful for querying. For example, for `WorkGraph<abc_name>` and `PythonJob<abc_name>`, the `node.label` will be `abc_name`.Note: `node.label` is mutable, thus user can modify it after.
---
 aiida_workgraph/calculations/python.py |  6 ++++++
 aiida_workgraph/engine/utils.py        |  2 +-
 aiida_workgraph/engine/workgraph.py    |  1 +
 tests/test_python.py                   | 15 +++++++++------
 tests/test_workgraph.py                |  1 +
 5 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/aiida_workgraph/calculations/python.py b/aiida_workgraph/calculations/python.py
index 7a65c563..8a469863 100644
--- a/aiida_workgraph/calculations/python.py
+++ b/aiida_workgraph/calculations/python.py
@@ -145,6 +145,12 @@ def _build_process_label(self) -> str:
         else:
             return f"PythonJob<{self.inputs.function_name.value}>"
 
+    def on_create(self) -> None:
+        """Called when a Process is created."""
+
+        super().on_create()
+        self.node.label = self.inputs.process_label.value
+
     def prepare_for_submission(self, folder: Folder) -> CalcInfo:
         """Prepare the calculation for submission.
 
diff --git a/aiida_workgraph/engine/utils.py b/aiida_workgraph/engine/utils.py
index ee091c17..6c2d5a95 100644
--- a/aiida_workgraph/engine/utils.py
+++ b/aiida_workgraph/engine/utils.py
@@ -95,7 +95,7 @@ def prepare_for_python_task(task: dict, kwargs: dict, var_kwargs: dict) -> dict:
     function_kwargs = serialize_to_aiida_nodes(function_kwargs)
     # transfer the args to kwargs
     inputs = {
-        "process_label": f"PythonJob<{task['name']}",
+        "process_label": f"PythonJob<{task['name']}>",
         "function_source_code": orm.Str(function_source_code),
         "function_name": orm.Str(function_name),
         "code": code,
diff --git a/aiida_workgraph/engine/workgraph.py b/aiida_workgraph/engine/workgraph.py
index 063dc5d7..91dc801b 100644
--- a/aiida_workgraph/engine/workgraph.py
+++ b/aiida_workgraph/engine/workgraph.py
@@ -416,6 +416,7 @@ def on_create(self) -> None:
         )
         saver = WorkGraphSaver(self.node, wgdata, restart_process=restart_process)
         saver.save()
+        self.node.label = wgdata["name"]
 
     def setup(self) -> None:
         # track if the awaitable callback is added to the runner
diff --git a/tests/test_python.py b/tests/test_python.py
index da09b150..d0a13906 100644
--- a/tests/test_python.py
+++ b/tests/test_python.py
@@ -24,23 +24,26 @@ def multiply(x: Any, y: Any) -> Any:
     wg = WorkGraph("test_PythonJob_outputs")
     wg.add_task(
         add,
-        name="add",
+        name="add1",
         x=1,
         y=2,
         computer="localhost",
     )
     wg.add_task(
         decorted_multiply,
-        name="multiply",
-        x=wg.tasks["add"].outputs["sum"],
+        name="multiply1",
+        x=wg.tasks["add1"].outputs["sum"],
         y=3,
         computer="localhost",
     )
     # wg.submit(wait=True)
     wg.run()
-    assert wg.tasks["add"].outputs["sum"].value.value == 3
-    assert wg.tasks["add"].outputs["diff"].value.value == -1
-    assert wg.tasks["multiply"].outputs["result"].value.value == 9
+    assert wg.tasks["add1"].outputs["sum"].value.value == 3
+    assert wg.tasks["add1"].outputs["diff"].value.value == -1
+    assert wg.tasks["multiply1"].outputs["result"].value.value == 9
+    # process_label and label
+    assert wg.tasks["add1"].node.process_label == "PythonJob<add1>"
+    assert wg.tasks["add1"].node.label == "add1"
 
 
 @pytest.mark.usefixtures("started_daemon_client")
diff --git a/tests/test_workgraph.py b/tests/test_workgraph.py
index e8ed049e..d2d6829f 100644
--- a/tests/test_workgraph.py
+++ b/tests/test_workgraph.py
@@ -39,6 +39,7 @@ def test_save_load(wg_calcfunction):
     wg.save()
     assert wg.process.process_state.value.upper() == "CREATED"
     assert wg.process.process_label == "WorkGraph<test_save_load>"
+    assert wg.process.label == "test_save_load"
     wg2 = WorkGraph.load(wg.process.pk)
     assert len(wg.tasks) == len(wg2.tasks)