Error handler can use the label of the ExitCode

aiida_workgraph/task.py

@@ -140,6 +140,7 @@ def error_handlers(self) -> list:
     def get_error_handlers(self) -> list:
         """Get the error handler function for this task."""
         from aiida_workgraph.utils import build_callable
+        from aiida.engine import ExitCode
 
         if self._error_handlers is None:
             return {}
@@ -152,6 +153,17 @@ def get_error_handlers(self) -> list:
             for handler in self._error_handlers:
                 handler["handler"] = build_callable(handler["handler"])
                 handlers[handler["handler"]["name"]] = handler
+        # convert exit code label (str) to status (int)
+        for handler in handlers.values():
+            exit_codes = []
+            for exit_code in handler["exit_codes"]:
+                if isinstance(exit_code, int):
+                    exit_codes.append(exit_code)
+                elif isinstance(exit_code, ExitCode):
+                    exit_codes.append(exit_code.status)
+                else:
+                    raise ValueError(f"Exit code {exit_code} is not a valid exit code.")
+            handler["exit_codes"] = exit_codes
         return handlers
 
     def _repr_mimebundle_(self, *args: Any, **kwargs: Any) -> any:

aiida_workgraph/workgraph.py

@@ -206,13 +206,28 @@ def to_dict(self, store_nodes=False) -> Dict[str, Any]:
 
     def get_error_handlers(self) -> Dict[str, Any]:
         """Get the error handlers."""
+        from aiida.engine import ExitCode
 
         from aiida_workgraph.utils import build_callable
 
         error_handlers = {}
         for name, error_handler in self._error_handlers.items():
             error_handler["handler"] = build_callable(error_handler["handler"])
             error_handlers[name] = error_handler
+        # convert exit code label (str) to status (int)
+        for handler in error_handlers.values():
+            for task in handler["tasks"].values():
+                exit_codes = []
+                for exit_code in task["exit_codes"]:
+                    if isinstance(exit_code, int):
+                        exit_codes.append(exit_code)
+                    elif isinstance(exit_code, ExitCode):
+                        exit_codes.append(exit_code.status)
+                    else:
+                        raise ValueError(
+                            f"Exit code {exit_code} is not a valid exit code."
+                        )
+                task["exit_codes"] = exit_codes
         return error_handlers
 
     def wait(self, timeout: int = 50, tasks: dict = None) -> None:

tests/test_error_handler.py

@@ -28,7 +28,15 @@ def handle_negative_sum(task: Task):
     wg.add_error_handler(
         handle_negative_sum,
         name="handle_negative_sum",
-        tasks={"add1": {"exit_codes": [410], "max_retries": 5, "kwargs": {}}},
+        tasks={
+            "add1": {
+                "exit_codes": [
+                    ArithmeticAddCalculation.exit_codes.ERROR_NEGATIVE_NUMBER
+                ],
+                "max_retries": 5,
+                "kwargs": {},
+            }
+        },
     )
     wg.submit(
         inputs={