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

Co-authored-by: Lin Zhihao <[email protected]>

Author: sitaowang1998

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,108 @@
+"""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
+
+
+# NOTE: This type alias is for clarification purposes only. It does not enforce static type checks.
+# Instead, we rely on the runtime check to ensure the first argument is `TaskContext`. To statically
+# enforce the first argument to be `TaskContext`, `Protocol` is required, which is not compatible
+# with `Callable` without explicit type casting.
+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 _validate_and_convert_params(signature: inspect.Signature) -> list[TaskInput]:
+    """
+    Validates the task parameters and converts them into a list of `core.TaskInput`.
+    :param signature:
+    :return: The converted task parameters.
+    :raises TypeError: If the parameters are invalid.
+    """
+    params = list(signature.parameters.values())
+    inputs = []
+    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 = "Parameters must have type annotation."
+            raise TypeError(msg)
+        tdl_type_str = to_tdl_type_str(param.annotation)
+        inputs.append(TaskInput(tdl_type_str, None))
+    return inputs
+
+
+def _validate_and_convert_return(signature: inspect.Signature) -> list[TaskOutput]:
+    """
+    Validates the task returns and converts them into a list of `core.TaskOutput`.
+    :param signature:
+    :return: The converted task returns.
+    :raises TypeError: If the return type is invalid.
+    """
+    returns = signature.return_annotation
+    outputs = []
+    if returns is inspect.Parameter.empty:
+        msg = "Return type must have type annotation."
+        raise TypeError(msg)
+
+    if not _is_tuple(returns):
+        tdl_type_str = to_tdl_type_str(returns)
+        if returns is Data:
+            outputs.append(TaskOutput(tdl_type_str, TaskOutputData()))
+        else:
+            outputs.append(TaskOutput(tdl_type_str, TaskOutputValue()))
+        return outputs
+
+    args = get_args(returns)
+    if Ellipsis in args:
+        msg = "Variable-length tuple return types are not supported."
+        raise TypeError(msg)
+    for arg in args:
+        tdl_type_str = to_tdl_type_str(arg)
+        if arg is Data:
+            outputs.append(TaskOutput(tdl_type_str, TaskOutputData()))
+        else:
+            outputs.append(TaskOutput(tdl_type_str, TaskOutputValue()))
+    return outputs
+
+
+def create_task(func: TaskFunction) -> core.Task:
+    """
+    Creates a core Task object from the task function.
+    :param func:
+    :return: The created core Task object.
+    :raise TypeError: If the function signature contains unsupported types.
+    """
+    if not isinstance(func, FunctionType):
+        msg = "`func` is not a function."
+        raise TypeError(msg)
+    signature = inspect.signature(func)
+    return core.Task(
+        function_name=func.__qualname__,
+        task_inputs=_validate_and_convert_params(signature),
+        task_outputs=_validate_and_convert_return(signature),
+    )

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

@@ -0,0 +1,67 @@
+"""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:
+    """
+    Represents a client-side task graph.
+
+    This class is a wrapper of `spider_py.core.Task`.
+    """
+
+    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: The new task graph.
+    :raises TypeError: If the parent outputs and child inputs do not match.
+    """
+    parent_core_graph: core.TaskGraph
+    child_core_graph: core.TaskGraph
+
+    if callable(parent):
+        parent_core_graph = core.TaskGraph()
+        parent_core_graph.add_task(create_task(parent))
+    else:
+        parent_core_graph = parent._impl
+
+    if callable(child):
+        child_core_graph = core.TaskGraph()
+        child_core_graph.add_task(create_task(child))
+    else:
+        child_core_graph = child._impl
+
+    graph = TaskGraph()
+    graph._impl = core.TaskGraph.chain_graph(parent_core_graph, child_core_graph)
+    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):