Refactor the syntax of decorator parameters for inputs, outputs and properties. (#123)

Refactor the syntax of decorator parameters for inputs, outputs and properties. Using a dictionary with keys and values to replace the list approach.

Author: superstar54

aiida_workgraph/collection.py

@@ -18,8 +18,8 @@ def new(
     ) -> Any:
         from aiida_workgraph.decorator import (
             build_task_from_callable,
-            build_python_task,
-            build_shell_task,
+            build_pythonjob_task,
+            build_shelljob_task,
         )
 
         # build the task on the fly if the identifier is a callable
@@ -31,13 +31,13 @@ def new(
                         "GraphBuilder task cannot be run remotely. Please set run_remotely=False."
                     )
                 # this is a PythonJob
-                identifier, _ = build_python_task(identifier)
+                identifier, _ = build_pythonjob_task(identifier)
             return super().new(identifier, name, uuid, **kwargs)
         if isinstance(identifier, str) and identifier.upper() == "PythonJob":
-            identifier, _ = build_python_task(kwargs.pop("function"))
+            identifier, _ = build_pythonjob_task(kwargs.pop("function"))
             return super().new(identifier, name, uuid, **kwargs)
         if isinstance(identifier, str) and identifier.upper() == "SHELLJOB":
-            identifier, _, links = build_shell_task(
+            identifier, _, links = build_shelljob_task(
                 nodes=kwargs.get("nodes", {}),
                 outputs=kwargs.get("outputs", None),
                 parser_outputs=kwargs.pop("parser_outputs", None),

aiida_workgraph/decorator.py

@@ -1,5 +1,5 @@
 from typing import Any, Callable, Dict, List, Optional, Union, Tuple
-from aiida_workgraph.utils import get_executor
+from aiida_workgraph.utils import get_executor, serialize_function
 from aiida.engine import calcfunction, workfunction, CalcJob, WorkChain
 from aiida import orm
 from aiida.orm.nodes.process.calculation.calcfunction import CalcFunctionNode
@@ -45,14 +45,18 @@ def add_input_recursive(
     else:
         port_name = f"{prefix}.{port.name}"
     required = port.required and required
-    input_names = [input[1] for input in inputs]
+    input_names = [input["name"] for input in inputs]
     if isinstance(port, PortNamespace):
         # TODO the default value is {} could cause problem, because the address of the dict is the same,
         # so if you change the value of one port, the value of all the ports of other tasks will be changed
         # consider to use None as default value
         if port_name not in input_names:
             inputs.append(
-                ["General", port_name, {"property": ["General", {"default": {}}]}]
+                {
+                    "identifier": "General",
+                    "name": port_name,
+                    "property": {"identifier": "General", "default": {}},
+                }
             )
         if required:
             args.append(port_name)
@@ -72,7 +76,7 @@ def add_input_recursive(
                 socket_type = aiida_socket_maping.get(port.valid_type[0], "General")
             else:
                 socket_type = aiida_socket_maping.get(port.valid_type, "General")
-            inputs.append([socket_type, port_name])
+            inputs.append({"identifier": socket_type, "name": port_name})
         if required:
             args.append(port_name)
         else:
@@ -92,18 +96,18 @@ def add_output_recursive(
     else:
         port_name = f"{prefix}.{port.name}"
     required = port.required and required
-    output_names = [output[1] for output in outputs]
+    output_names = [output["name"] for output in outputs]
     if isinstance(port, PortNamespace):
         # TODO the default value is {} could cause problem, because the address of the dict is the same,
         # so if you change the value of one port, the value of all the ports of other tasks will be changed
         # consider to use None as default value
         if port_name not in output_names:
-            outputs.append(["General", port_name])
+            outputs.append({"identifier": "General", "name": port_name})
         for value in port.values():
             add_output_recursive(outputs, value, prefix=port_name, required=required)
     else:
         if port_name not in output_names:
-            outputs.append(["General", port_name])
+            outputs.append({"identifier": "General", "name": port_name})
     return outputs
 
 
@@ -213,35 +217,45 @@ def build_task_from_AiiDA(
                 or executor.process_class._var_positional
             )
         tdata["var_kwargs"] = name
-        inputs.append(["General", name, {"property": ["General", {"default": {}}]}])
+        inputs.append(
+            {
+                "identifier": "General",
+                "name": name,
+                "property": {"identifier": "General", "default": {}},
+            }
+        )
+    # TODO In order to reload the WorkGraph from process, "is_pickle" should be True
+    # so I pickled the function here, but this is not necessary
+    # we need to update the node_graph to support the path and name of the function
+    tdata["identifier"] = tdata.pop("identifier", tdata["executor"].__name__)
+    tdata["executor"] = {
+        "executor": pickle.dumps(executor),
+        "type": tdata["task_type"],
+        "is_pickle": True,
+    }
     if tdata["task_type"].upper() in ["CALCFUNCTION", "WORKFUNCTION"]:
-        outputs = [["General", "result"]] if not outputs else outputs
+        outputs = (
+            [{"identifier": "General", "name": "result"}] if not outputs else outputs
+        )
+        # get the source code of the function
+        tdata["executor"] = serialize_function(executor)
+        # tdata["executor"]["type"] = tdata["task_type"]
     # print("kwargs: ", kwargs)
     # add built-in sockets
-    outputs.append(["General", "_outputs"])
-    outputs.append(["General", "_wait"])
-    inputs.append(["General", "_wait", {"link_limit": 1e6}])
+    outputs.append({"identifier": "General", "name": "_outputs"})
+    outputs.append({"identifier": "General", "name": "_wait"})
+    inputs.append({"identifier": "General", "name": "_wait", "link_limit": 1e6})
     tdata["node_class"] = Task
     tdata["args"] = args
     tdata["kwargs"] = kwargs
     tdata["inputs"] = inputs
     tdata["outputs"] = outputs
-    tdata["identifier"] = tdata.pop("identifier", tdata["executor"].__name__)
-    # TODO In order to reload the WorkGraph from process, "is_pickle" should be True
-    # so I pickled the function here, but this is not necessary
-    # we need to update the node_graph to support the path and name of the function
-    executor = {
-        "executor": pickle.dumps(executor),
-        "type": tdata["task_type"],
-        "is_pickle": True,
-    }
-    tdata["executor"] = executor
     task = create_task(tdata)
     task.is_aiida_component = True
     return task, tdata
 
 
-def build_python_task(func: Callable) -> Task:
+def build_pythonjob_task(func: Callable) -> Task:
     """Build PythonJob task from function."""
     from aiida_workgraph.calculations.python import PythonJob
     from copy import deepcopy
@@ -254,10 +268,10 @@ def build_python_task(func: Callable) -> Task:
     inputs = tdata["inputs"]
     inputs.extend(
         [
-            ["String", "computer"],
-            ["String", "code_label"],
-            ["String", "code_path"],
-            ["String", "prepend_text"],
+            {"identifier": "String", "name": "computer"},
+            {"identifier": "String", "name": "code_label"},
+            {"identifier": "String", "name": "code_path"},
+            {"identifier": "String", "name": "prepend_text"},
         ]
     )
     outputs = tdata["outputs"]
@@ -269,8 +283,8 @@ def build_python_task(func: Callable) -> Task:
             outputs.append(output)
     # change "copy_files" link_limit to 1e6
     for input in inputs:
-        if input[1] == "copy_files":
-            input[2].update({"link_limit": 1e6})
+        if input["name"] == "copy_files":
+            input["link_limit"] = 1e6
     # append the kwargs of the PythonJob task to the function task
     kwargs = tdata["kwargs"]
     kwargs.extend(["computer", "code_label", "code_path", "prepend_text"])
@@ -284,7 +298,7 @@ def build_python_task(func: Callable) -> Task:
     return task, tdata
 
 
-def build_shell_task(
+def build_shelljob_task(
     nodes: dict = None, outputs: list = None, parser_outputs: list = None
 ) -> Task:
     """Build ShellJob with custom inputs and outputs."""
@@ -300,7 +314,7 @@ def build_shell_task(
     nodes = {} if nodes is None else nodes
     keys = list(nodes.keys())
     for key in keys:
-        inputs.append(["General", f"nodes.{key}"])
+        inputs.append({"identifier": "General", "name": f"nodes.{key}"})
         # input is a output of another task, we make a link
         if isinstance(nodes[key], NodeSocket):
             links[f"nodes.{key}"] = nodes[key]
@@ -309,12 +323,20 @@ def build_shell_task(
     for input in inputs:
         if input not in tdata["inputs"]:
             tdata["inputs"].append(input)
-            tdata["kwargs"].append(input[1])
+            tdata["kwargs"].append(input["name"])
     # Extend the outputs
-    tdata["outputs"].extend([["General", "stdout"], ["General", "stderr"]])
+    tdata["outputs"].extend(
+        [
+            {"identifier": "General", "name": "stdout"},
+            {"identifier": "General", "name": "stderr"},
+        ]
+    )
     outputs = [] if outputs is None else outputs
     parser_outputs = [] if parser_outputs is None else parser_outputs
-    outputs = [["General", ShellParser.format_link_label(output)] for output in outputs]
+    outputs = [
+        {"identifier": "General", "name": ShellParser.format_link_label(output)}
+        for output in outputs
+    ]
     outputs.extend(parser_outputs)
     # add user defined outputs
     for output in outputs:
@@ -324,8 +346,8 @@ def build_shell_task(
     tdata["identifier"] = "ShellJob"
     tdata["inputs"].extend(
         [
-            ["General", "command"],
-            ["General", "resolve_command"],
+            {"identifier": "General", "name": "command"},
+            {"identifier": "General", "name": "resolve_command"},
         ]
     )
     tdata["kwargs"].extend(["command", "resolve_command"])
@@ -346,25 +368,29 @@ def build_task_from_workgraph(wg: any) -> Task:
     # add all the inputs/outputs from the tasks in the workgraph
     for task in wg.tasks:
         # inputs
-        inputs.append(["General", f"{task.name}"])
+        inputs.append({"identifier": "General", "name": f"{task.name}"})
         for socket in task.inputs:
             if socket.name == "_wait":
                 continue
-            inputs.append(["General", f"{task.name}.{socket.name}"])
+            inputs.append(
+                {"identifier": "General", "name": f"{task.name}.{socket.name}"}
+            )
         # outputs
-        outputs.append(["General", f"{task.name}"])
+        outputs.append({"identifier": "General", "name": f"{task.name}"})
         for socket in task.outputs:
             if socket.name in ["_wait", "_outputs"]:
                 continue
-            outputs.append(["General", f"{task.name}.{socket.name}"])
+            outputs.append(
+                {"identifier": "General", "name": f"{task.name}.{socket.name}"}
+            )
             group_outputs.append(
                 [f"{task.name}.{socket.name}", f"{task.name}.{socket.name}"]
             )
-    kwargs = [input[1] for input in inputs]
+    kwargs = [input["name"] for input in inputs]
     # add built-in sockets
-    outputs.append(["General", "_outputs"])
-    outputs.append(["General", "_wait"])
-    inputs.append(["General", "_wait", {"link_limit": 1e6}])
+    outputs.append({"identifier": "General", "name": "_outputs"})
+    outputs.append({"identifier": "General", "name": "_wait"})
+    inputs.append({"identifier": "General", "name": "_wait", "link_limit": 1e6})
     tdata["node_class"] = Task
     tdata["kwargs"] = kwargs
     tdata["inputs"] = inputs
@@ -385,77 +411,6 @@ def build_task_from_workgraph(wg: any) -> Task:
     return task
 
 
-def get_required_imports(func):
-    """Retrieve type hints and the corresponding module"""
-    from typing import get_type_hints, _SpecialForm
-
-    type_hints = get_type_hints(func)
-    imports = {}
-
-    def add_imports(type_hint):
-        if isinstance(
-            type_hint, _SpecialForm
-        ):  # Handle special forms like Any, Union, Optional
-            module_name = "typing"
-            type_name = type_hint._name or str(type_hint)
-        elif hasattr(
-            type_hint, "__origin__"
-        ):  # This checks for higher-order types like List, Dict
-            module_name = type_hint.__module__
-            type_name = type_hint._name
-            for arg in type_hint.__args__:
-                if arg is type(None):  # noqa: E721
-                    continue
-                add_imports(arg)  # Recursively add imports for each argument
-        elif hasattr(type_hint, "__module__"):
-            module_name = type_hint.__module__
-            type_name = type_hint.__name__
-        else:
-            return  # If no module or origin, we can't import it, e.g., for literals
-
-        if module_name not in imports:
-            imports[module_name] = set()
-        imports[module_name].add(type_name)
-
-    for _, type_hint in type_hints.items():
-        add_imports(type_hint)
-
-    return imports
-
-
-def serialize_function(func: Callable) -> Dict[str, Any]:
-    """Serialize a function for storage or transmission."""
-    import cloudpickle as pickle
-    import inspect
-    import textwrap
-
-    source_code = inspect.getsource(func)
-    source_code_lines = source_code.split("\n")
-    # we need save the source code explicitly, because in the case of jupyter notebook,
-    # the source code is not saved in the pickle file
-    function_source_code = "\n".join(source_code_lines[1:])
-    function_source_code = textwrap.dedent(function_source_code)
-    # we also need to include the necessary imports for the types used in the type hints.
-    try:
-        required_imports = get_required_imports(func)
-    except Exception as e:
-        required_imports = {}
-        print(f"Failed to get required imports for function {func.__name__}: {e}")
-    # Generate import statements
-    import_statements = "\n".join(
-        f"from {module} import {', '.join(types)}"
-        for module, types in required_imports.items()
-    )
-    return {
-        "executor": pickle.dumps(func),
-        "type": "function",
-        "is_pickle": True,
-        "function_name": func.__name__,
-        "function_source_code": function_source_code,
-        "import_statements": import_statements,
-    }
-
-
 def generate_tdata(
     func: Callable,
     identifier: str,
@@ -475,9 +430,9 @@ def generate_tdata(
     )
     task_outputs = outputs
     # add built-in sockets
-    _inputs.append(["General", "_wait", {"link_limit": 1e6}])
-    task_outputs.append(["General", "_wait"])
-    task_outputs.append(["General", "_outputs"])
+    _inputs.append({"identifier": "General", "name": "_wait", "link_limit": 1e6})
+    task_outputs.append({"identifier": "General", "name": "_wait"})
+    task_outputs.append({"identifier": "General", "name": "_outputs"})
     tdata = {
         "node_class": Task,
         "identifier": identifier,
@@ -536,7 +491,7 @@ def decorator(func):
                 func,
                 identifier,
                 inputs or [],
-                outputs or [["General", "result"]],
+                outputs or [{"identifier": "General", "name": "result"}],
                 properties or [],
                 catalog,
                 task_type,
@@ -577,7 +532,9 @@ def decorator(func):
             # use cloudpickle to serialize function
             func.identifier = identifier
 
-            task_outputs = [["General", output[1]] for output in outputs]
+            task_outputs = [
+                {"identifier": "General", "name": output[1]} for output in outputs
+            ]
             # print(task_inputs, task_outputs)
             #
             task_type = "graph_builder"

aiida_workgraph/engine/utils.py

@@ -86,7 +86,7 @@ def prepare_for_python_task(task: dict, kwargs: dict, var_kwargs: dict) -> dict:
     function_source_code = (
         task["executor"]["import_statements"]
         + "\n"
-        + task["executor"]["function_source_code"]
+        + task["executor"]["function_source_code_without_decorator"]
     )
     # outputs
     output_name_list = [output["name"] for output in task["outputs"]]