feat(spider-py): Add support for client-end task graph grouping and chaining.

python/spider-py/pyproject.toml

@@ -74,10 +74,14 @@ typing-extensions = false  # Disable type extensions for 3.10 compatibility
 [tool.ruff.lint.per-file-ignores]
 "tests/**" = [
   "INP001",  # Allow implicit namespace package for tests
+  "PLR2004",  # Allow use of magic value
   "S101",  # Allow use of `assert` (security warning)
   "S603",  # Allow use of `subprocess.Popen` (security warning)
   "T201",  # Allow use of `print` (testing)
 ]
 
+[tool.ruff.lint.flake8-self]
+ignore-names = ["_impl"]
+
 [tool.ruff.lint.pydocstyle]
 ignore-decorators = ["typing.override"]

python/spider-py/src/spider_py/__init__.py

@@ -1,5 +1,6 @@
 """Spider package root."""
 
+from spider_py.client import chain, group, TaskContext, TaskGraph
 from spider_py.type import Double, Float, Int8, Int16, Int32, Int64
 
 __all__ = [
@@ -9,4 +10,8 @@
     "Int16",
     "Int32",
     "Int64",
+    "TaskContext",
+    "TaskGraph",
+    "chain",
+    "group",
 ]

python/spider-py/src/spider_py/client/__init__.py

@@ -1 +1,11 @@
 """Spider python client."""
+
+from .task import TaskContext
+from .task_graph import chain, group, TaskGraph
+
+__all__ = [
+    "TaskContext",
+    "TaskGraph",
+    "chain",
+    "group",
+]

python/spider-py/src/spider_py/client/data.py

@@ -0,0 +1,5 @@
+"""Spider client Data module."""
+
+
+class Data:
+    """Represents a spider client data."""

python/spider-py/src/spider_py/client/task.py

@@ -0,0 +1,101 @@
+"""Spider client task module."""
+
+import inspect
+from collections.abc import Callable
+from types import FunctionType, GenericAlias
+from typing import get_args, get_origin
+
+from spider_py import core
+from spider_py.client.data import Data
+from spider_py.core import TaskInput, TaskOutput, TaskOutputData, TaskOutputValue
+from spider_py.type import to_tdl_type_str
+
+
+class TaskContext:
+    """Spider task context."""
+
+    # TODO: Implement task context for use in task executor
+
+
+# Check the TaskFunction signature at runtime.
+# Enforcing static check for first argument requires the use of Protocol. However, functions, which
+# are Callable, are not considered a Protocol without explicit cast.
+TaskFunction = Callable[..., object]
+
+
+def _is_tuple(t: type | GenericAlias) -> bool:
+    """
+    :param t:
+    :return: Whether t is a tuple.
+    """
+    return get_origin(t) is tuple
+
+
+def _process_parameters(task: core.Task, signature: inspect.Signature) -> None:
+    """
+    Checks the parameters validity and add them to the task.
+    :param task:
+    :param signature:
+    :raises TypeError: If the parameters are invalid.
+    """
+    params = list(signature.parameters.values())
+    if not params or params[0].annotation is not TaskContext:
+        msg = "First argument is not a TaskContext."
+        raise TypeError(msg)
+    for param in params[1:]:
+        if param.kind in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}:
+            msg = "Variadic parameters are not supported."
+            raise TypeError(msg)
+        if param.annotation is inspect.Parameter.empty:
+            msg = "Argument must have type annotation"
+            raise TypeError(msg)
+        tdl_type_str = to_tdl_type_str(param.annotation)
+        task.task_inputs.append(TaskInput(tdl_type_str, None))
+
+
+def _process_return(task: core.Task, signature: inspect.Signature) -> None:
+    """
+    Checks the return type validity and add them to the task.
+    :param task:
+    :param signature:
+    :raises TypeError: If the return type is invalid.
+    """
+    returns = signature.return_annotation
+    if returns is inspect.Parameter.empty:
+        msg = "Return type must have type annotation"
+        raise TypeError(msg)
+    if _is_tuple(returns):
+        args = get_args(returns)
+        if Ellipsis in args:
+            msg = "Variable-length tuple return types are not supported."
+            raise TypeError(msg)
+        for r in args:
+            tdl_type_str = to_tdl_type_str(r)
+            if r is Data:
+                task.task_outputs.append(TaskOutput(tdl_type_str, TaskOutputData()))
+            else:
+                task.task_outputs.append(TaskOutput(tdl_type_str, TaskOutputValue()))
+    else:
+        tdl_type_str = to_tdl_type_str(returns)
+        if returns is Data:
+            task.task_outputs.append(TaskOutput(tdl_type_str, TaskOutputData()))
+        else:
+            task.task_outputs.append(TaskOutput(tdl_type_str, TaskOutputValue()))
+
+
+def _create_task(func: TaskFunction) -> core.Task:
+    """
+    Creates a core Task object from the task function.
+    :param func:
+    :return:
+    :raise TypeError: If the function signature contains unsupported types.
+    """
+    task = core.Task()
+    if not isinstance(func, FunctionType):
+        msg = "`func` is not a function."
+        raise TypeError(msg)
+    task.function_name = func.__qualname__
+    signature = inspect.signature(func)
+    _process_parameters(task, signature)
+    _process_return(task, signature)
+    return task

python/spider-py/src/spider_py/client/task_graph.py

@@ -0,0 +1,59 @@
+"""Spider client TaskGraph module."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from spider_py import core
+from spider_py.client.task import _create_task, TaskFunction
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+class TaskGraph:
+    """
+    Spider client TaskGraph class.
+    Wraps around the core TaskGraph class.
+    """
+
+    def __init__(self) -> None:
+        """Initializes TaskGraph."""
+        self._impl = core.TaskGraph()
+
+
+def group(tasks: Sequence[TaskFunction | TaskGraph]) -> TaskGraph:
+    """
+    Groups task functions and task graph into a single task graph.
+    :param tasks: List of task functions or task graphs.
+    :return: The new task graph.
+    """
+    graph = TaskGraph()
+    for task in tasks:
+        if callable(task):
+            graph._impl.add_task(_create_task(task))
+        else:
+            graph._impl.merge_graph(task._impl)
+
+    return graph
+
+
+def chain(parent: TaskFunction | TaskGraph, child: TaskFunction | TaskGraph) -> TaskGraph:
+    """
+    Chains two task functions or task graphs into a single task graph.
+    :param parent:
+    :param child:
+    :return:
+    :raises TypeError: If the parent outputs and child inputs do not match.
+    """
+    if callable(parent):
+        task = _create_task(parent)
+        parent = TaskGraph()
+        parent._impl.add_task(task)
+    if callable(child):
+        task = _create_task(child)
+        child = TaskGraph()
+        child._impl.add_task(task)
+    graph = TaskGraph()
+    graph._impl = core.TaskGraph.chain_graph(parent._impl, child._impl)
+    return graph

python/spider-py/src/spider_py/core/task.py

@@ -19,11 +19,26 @@ class TaskInputOutput:
 
 TaskInputValue = bytes
 TaskInputData = DataId
-TaskInput = TaskInputOutput | TaskInputValue | TaskInputData
+
+
+@dataclass
+class TaskInput:
+    """Represents a task input"""
+
+    type: str
+    value: TaskInputData | TaskInputOutput | None
+
 
 TaskOutputValue = bytes
 TaskOutputData = DataId
-TaskOutput = TaskOutputValue | TaskOutputData
+
+
+@dataclass
+class TaskOutput:
+    """Represents a task output"""
+
+    type: str
+    value: TaskOutputData | TaskOutputValue
 
 
 class TaskState(IntEnum):