bitool 0.1.2__py3-none-any.whl

bitool/utils/task_group.py

@@ -0,0 +1,668 @@
+"""任务编排模块 - 用于定义和管理任务流.
+
+本模块提供通用的任务组数据结构,允许通过简单的配置
+来定义复杂的命令执行流程.
+
+核心功能
+--------
+- 声明式任务组定义
+- 支持命令列表和可调用对象
+- 支持顺序执行和并行执行
+- 灵活的条件执行控制
+- 任务依赖管理 (DAG调度)
+
+快速开始
+--------
+>>> from pytola_core.command.task_group import TaskGroup, Task
+>>> group = TaskGroup(
+...     name="git-init",
+...     commands=[
+...         Task(cmd=["git", "init"]),
+...         Task(cmd=["git", "add", "."]),
+...         Task(cmd=["git", "commit", "-m", "initial commit"]),
+...     ]
+... )
+>>> group.execute()
+
+DAG 依赖调度
+-----------
+>>> # 定义带依赖关系的任务流
+>>> group = TaskGroup(
+...     name="build-pipeline",
+...     commands=[
+...         Task(name="init", cmd=["git", "init"]),
+...         Task(name="install", cmd=["pip", "install", "-r", "requirements.txt"],
+...                    depends_on=["init"]),
+...         Task(name="lint", cmd=["ruff", "check", "."], depends_on=["install"]),
+...         Task(name="test", cmd=["pytest"], depends_on=["install"]),
+...         Task(name="build", cmd=["python", "-m", "build"],
+...                    depends_on=["lint", "test"]),
+...     ],
+...     max_workers=4,
+... )
+>>> # lint 和 test 将并行执行(都依赖 install,都被 build 依赖)
+>>> group.execute()
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+from bitool.utils import execute
+
+from ..core import logger
+from .executor import RunResult
+from .task import Task
+
+
+@dataclass
+class TaskGroup:
+    """任务组 - 用于组织和管理相关任务.
+
+    Attributes
+    ----------
+    name : str
+        任务组名称
+    commands : list[Task]
+        任务项列表
+    description : str
+        任务组描述
+    enabled : bool
+        是否启用此任务组
+    parallel : bool
+        是否启用并行执行,默认 False (顺序执行)
+    max_workers : int | None
+        最大并行工作线程数,None 表示使用默认值
+    stop_on_failure : bool
+        失败时是否停止执行,默认 False
+    """
+
+    name: str
+    commands: list[Task] = field(default_factory=list)
+    description: str = ""
+    enabled: bool = True
+    parallel: bool = False
+    max_workers: int | None = None
+    stop_on_failure: bool = False
+
+    def add_command(
+        self,
+        cmd: list[str] | None = None,
+        func: Callable[[], RunResult] | None = None,
+        condition: Callable[[], bool] | None = None,
+        description: str = "",
+        *,
+        parallel: bool = False,
+        name: str = "",
+        depends_on: list[str] | None = None,
+    ) -> TaskGroup:
+        """添加任务项.
+
+        Parameters
+        ----------
+        cmd : list[str], optional
+            命令及其参数列表
+        func : Callable[[], RunResult], optional
+            可执行的函数
+        condition : Callable[[], bool], optional
+            执行条件函数
+        description : str
+            任务描述
+        parallel : bool
+            是否允许并行执行
+        name : str
+            任务唯一标识,用于依赖引用
+        depends_on : list[str], optional
+            依赖的任务名称列表
+
+        Returns
+        -------
+        TaskGroup
+            返回自身以支持链式调用
+        """
+        if depends_on is None:
+            depends_on = []
+        self.commands.append(
+            Task(
+                cmd=cmd,
+                func=func,
+                condition=condition,
+                description=description,
+                parallel=parallel,
+                name=name,
+                depends_on=depends_on,
+            )
+        )
+        return self
+
+    def add_commands(self, commands: list[Task]) -> TaskGroup:
+        """批量添加任务项.
+
+        Parameters
+        ----------
+        commands : list[Task]
+            任务项列表
+
+        Returns
+        -------
+        TaskGroup
+            返回自身以支持链式调用
+        """
+        self.commands.extend(commands)
+        return self
+
+    def execute(self) -> list[RunResult]:
+        """执行任务组.
+
+        Returns
+        -------
+        list[RunResult]
+            任务执行结果列表
+        """
+        if not self.enabled:
+            logger.info(f"任务组 '{self.name}' 已禁用,跳过执行")
+            return []
+
+        if not self.commands:
+            logger.warning(f"任务组 '{self.name}' 没有可执行的任务")
+            return []
+
+        logger.info(f"开始执行任务组: '{self.name}' ({len(self.commands)} 个任务)")
+
+        # 过滤需要执行的任务
+        executable_commands = []
+        for item in self.commands:
+            if item.is_executable:
+                exec_item = item.executable_cmd
+                if exec_item is not None:
+                    executable_commands.append((item, exec_item))
+            else:
+                logger.debug(
+                    f"跳过任务 (条件不满足): {item.description or item.cmd or item.func}"
+                )
+
+        if not executable_commands:
+            logger.info(f"任务组 '{self.name}' 没有满足条件的任务")
+            return []
+
+        # 智能选择执行策略
+        has_dependencies = any(item.depends_on for item, _ in executable_commands)
+
+        if has_dependencies:
+            # 使用 DAG 调度
+            logger.info(f"任务组 '{self.name}' 使用 DAG 依赖调度")
+            results = self._execute_with_dag(executable_commands)
+        elif self.parallel:
+            # 使用原有并行逻辑
+            results = self._execute_parallel(executable_commands)
+        else:
+            # 使用原有串行逻辑
+            results = self._execute_sequential(executable_commands)
+
+        # 统计执行结果
+        success_count = sum(1 for r in results if r.success)
+        total_count = len(results)
+
+        logger.info(
+            f"任务组 '{self.name}' 执行完成 | "
+            f"总数:{total_count} | 成功:{success_count} | 失败:{total_count - success_count}"
+        )
+
+        return results
+
+    def _execute_sequential(
+        self, commands: list[tuple[Task, list[str] | Callable[[], RunResult]]]
+    ) -> list[RunResult]:
+        """顺序执行任务.
+
+        Parameters
+        ----------
+        commands : list[tuple[Task, list[str] | Callable[[], RunResult]]]
+            任务项和可执行对象的元组列表
+
+        Returns
+        -------
+        list[RunResult]
+            任务执行结果列表
+        """
+        results = []
+        for item, _executable in commands:
+            # 优先使用 description，其次使用命令列表，最后使用函数名
+            if item.description:
+                task_desc = item.description
+            elif item.cmd:
+                task_desc = " ".join(item.cmd)
+            elif item.func and hasattr(item.func, "__name__"):
+                task_desc = item.func.__name__
+            else:
+                task_desc = str(item.func) if item.func else "unknown"
+
+            logger.info(f"执行: `{task_desc}`")
+
+            # 执行单个任务：通过 cmd/func 属性明确区分任务类型
+            if item.cmd is not None:
+                result = execute(item.cmd)
+            elif item.func is not None:
+                result = item.func()
+            else:
+                msg = "任务未定义 cmd 或 func"
+                raise ValueError(msg)
+            results.append(result)
+
+            # 检查是否需要停止
+            if not result.success and self.stop_on_failure:
+                logger.warning(f"任务失败,停止执行: `{task_desc}`")
+                break
+
+        return results
+
+    def _find_command(
+        self,
+        commands: list[tuple[Task, list[str] | Callable[[], RunResult]]],
+        identifier: str | int,
+    ) -> tuple[Task, list[str] | Callable[[], RunResult]]:
+        """根据标识查找任务.
+
+        Parameters
+        ----------
+        commands : list[tuple[Task, list[str] | Callable[[], RunResult]]]
+            任务项和可执行对象的元组列表
+        identifier : str | int
+            任务名称或ID
+
+        Returns
+        -------
+        tuple[Task, list[str] | Callable[[], RunResult]]
+            任务项和可执行对象的元组
+        """
+        for item, executable in commands:
+            if (item.name or id(item)) == identifier:
+                return item, executable
+        msg = f"找不到任务: {identifier}"
+        raise ValueError(msg)
+
+    def _update_dependents(
+        self,
+        name: str | int,
+        graph: dict[str | int, list[str | int]],
+        in_degree: dict[str | int, int],
+        ready_queue: list[str | int],
+    ) -> None:
+        """更新依赖节点,将入度为0的加入就绪队列.
+
+        Parameters
+        ----------
+        name : str | int
+            已完成的任务ID
+        graph : dict[str | int, list[str | int]]
+            依赖图邻接表
+        in_degree : dict[str | int, int]
+            各节点入度
+        ready_queue : list[str | int]
+            就绪队列
+        """
+        for dependent_id in graph.get(name, []):
+            in_degree[dependent_id] -= 1
+            if in_degree[dependent_id] == 0:
+                ready_queue.append(dependent_id)
+
+    def _build_dependency_graph(
+        self, commands: list[tuple[Task, list[str] | Callable[[], RunResult]]]
+    ) -> tuple[dict[str | int, list[str | int]], dict[str | int, int]]:
+        """构建依赖图,检测循环依赖.
+
+        Parameters
+        ----------
+        commands : list[tuple[Task, list[str] | Callable[[], RunResult]]]
+            任务项和可执行对象的元组列表
+
+        Returns
+        -------
+        tuple[dict[str | int, list[str | int]], dict[str | int, int]]
+            依赖图邻接表和入度表
+
+        Raises
+        ------
+        ValueError
+            如果存在循环依赖或引用不存在的任务名称
+        """
+        graph: dict[str | int, list[str | int]] = {}
+        in_degree: dict[str | int, int] = {}
+
+        # 初始化图
+        for item, _ in commands:
+            name = item.name or id(item)
+            graph[name] = []
+            in_degree[name] = len(item.depends_on)
+
+        # 建立邻接表
+        for item, _ in commands:
+            name = item.name or id(item)
+            for dep_id in item.depends_on:
+                # 查找依赖的ID
+                dep_found = False
+                for dep_item, _ in commands:
+                    if dep_item.name == dep_id:
+                        dep_found = True
+                        graph[dep_id].append(name)
+                        break
+
+                if not dep_found:
+                    msg = f"依赖的任务名称不存在: `{dep_id}`"
+                    raise ValueError(msg)
+
+        # 检测循环依赖
+        if self._detect_cycle(graph):
+            msg = "检测到循环依赖,请检查任务的 depends_on 配置"
+            raise ValueError(msg)
+
+        return graph, in_degree
+
+    def _detect_cycle(self, graph: dict[str | int, list[str | int]]) -> bool:
+        """检测循环依赖.
+
+        Parameters
+        ----------
+        graph : dict[str | int, list[str | int]]
+            依赖图邻接表
+
+        Returns
+        -------
+        bool
+            如果存在循环依赖返回True
+        """
+        visited: set[str | int] = set()
+        rec_stack: set[str | int] = set()
+
+        def dfs(node: str | int) -> bool:
+            visited.add(node)
+            rec_stack.add(node)
+
+            for neighbor in graph.get(node, []):
+                if neighbor not in visited:
+                    if dfs(neighbor):
+                        return True
+                elif neighbor in rec_stack:
+                    return True
+
+            rec_stack.remove(node)
+            return False
+
+        return any(node not in visited and dfs(node) for node in graph)
+
+    def _execute_parallel(
+        self, commands: list[tuple[Task, list[str] | Callable[[], RunResult]]]
+    ) -> list[RunResult]:
+        """并行执行任务.
+
+        Parameters
+        ----------
+        commands : list[tuple[Task, list[str] | Callable[[], RunResult]]]
+            任务项和可执行对象的元组列表
+
+        Returns
+        -------
+        list[RunResult]
+            任务执行结果列表
+        """
+        from concurrent.futures import ThreadPoolExecutor, as_completed
+
+        # 分离顺序和并行任务
+        sequential_commands = [
+            (item, exec_cmd) for item, exec_cmd in commands if not item.parallel
+        ]
+        parallel_commands = [
+            (item, exec_cmd) for item, exec_cmd in commands if item.parallel
+        ]
+
+        results = []
+
+        # 先执行顺序任务
+        if sequential_commands:
+            logger.info(f"顺序执行 {len(sequential_commands)} 个任务")
+            seq_results = self._execute_sequential(sequential_commands)
+            results.extend(seq_results)
+
+            # 如果失败且设置了停止标志,则不执行并行任务
+            if self.stop_on_failure and any(not r.success for r in seq_results):
+                logger.warning("顺序任务失败,跳过并行任务")
+                return results
+
+        # 并行执行任务
+        if parallel_commands:
+            logger.info(f"并行执行 {len(parallel_commands)} 个任务")
+            max_workers = self.max_workers if self.max_workers is not None else 4
+
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                future_to_item = {}
+                for item, _ in parallel_commands:
+                    # 通过 cmd/func 属性明确区分任务类型
+                    if item.cmd is not None:
+                        from bitool.utils.executor import execute
+
+                        future = executor.submit(execute, item.cmd)
+                    elif item.func is not None:
+                        future = executor.submit(item.func)
+                    else:
+                        continue
+                    future_to_item[future] = item
+
+                # 收集结果
+                for future in as_completed(future_to_item):
+                    item = future_to_item[future]
+                    try:
+                        result = future.result()
+                        results.append(result)
+                        # 优先使用 description，其次使用命令列表，最后使用函数名
+                        if item.description:
+                            task_desc = item.description
+                        elif item.cmd:
+                            task_desc = " ".join(item.cmd)
+                        elif item.func and hasattr(item.func, "__name__"):
+                            task_desc = item.func.__name__
+                        else:
+                            task_desc = str(item.func) if item.func else "unknown"
+                        logger.info(f"完成: {task_desc}")
+                    except (OSError, RuntimeError, ValueError) as e:
+                        if item.description:
+                            task_desc = item.description
+                        elif item.cmd:
+                            task_desc = " ".join(item.cmd)
+                        elif item.func and hasattr(item.func, "__name__"):
+                            task_desc = item.func.__name__
+                        else:
+                            task_desc = str(item.func) if item.func else "unknown"
+                        logger.error(f"并行任务异常: {task_desc}, 错误:{e}")
+                        results.append(
+                            RunResult(
+                                returncode=-1,
+                                stdout="",
+                                stderr=str(e),
+                                execution_time=0.0,
+                                mode="failed",
+                                cmd=item.cmd if item.cmd else [],
+                            )
+                        )
+
+        return results
+
+    def _execute_with_dag(
+        self, commands: list[tuple[Task, list[str] | Callable[[], RunResult]]]
+    ) -> list[RunResult]:
+        """基于 DAG 的智能调度执行.
+
+        根据任务间的依赖关系,自动识别可并行执行的任务,
+        并按依赖顺序调度执行.
+
+        Parameters
+        ----------
+        commands : list[tuple[Task, list[str] | Callable[[], RunResult]]]
+            任务项和可执行对象的元组列表
+
+        Returns
+        -------
+        list[RunResult]
+            任务执行结果列表
+        """
+        from concurrent.futures import ThreadPoolExecutor, as_completed
+
+        graph, in_degree = self._build_dependency_graph(commands)
+
+        # 找到所有入度为0的节点(无依赖)
+        ready_queue: list[str | int] = [
+            name for name, degree in in_degree.items() if degree == 0
+        ]
+
+        results: dict[str | int, RunResult] = {}
+        stopped = False
+
+        while ready_queue and not stopped:
+            # 当前批次可以并行执行的任务
+            current_batch = ready_queue[:]
+            ready_queue.clear()
+
+            # 并行执行当前批次
+            max_workers = (
+                self.max_workers if self.max_workers is not None else len(current_batch)
+            )
+
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                future_to_cmd: dict[Any, str | int] = {}
+
+                for name in current_batch:
+                    item, _ = self._find_command(commands, name)
+
+                    # 检查条件
+                    if not item.is_executable:
+                        logger.debug(
+                            f"跳过任务 (条件不满足): {item.description or name}"
+                        )
+                        # 标记为已完成,更新依赖
+                        self._update_dependents(name, graph, in_degree, ready_queue)
+                        continue
+
+                    logger.info(f"DAG调度执行: {item.description or name}")
+
+                    # 提交任务
+                    if item.cmd is not None:
+                        from bitool.utils.executor import execute
+
+                        future = executor.submit(execute, item.cmd)
+                    elif item.func is not None:
+                        future = executor.submit(item.func)
+                    else:
+                        continue
+
+                    future_to_cmd[future] = name
+
+                # 收集结果
+                for future in as_completed(future_to_cmd):
+                    name = future_to_cmd[future]
+                    item, _ = self._find_command(commands, name)
+
+                    try:
+                        result = future.result()
+                        results[name] = result
+                        task_desc = item.description or name
+                        logger.info(f"DAG完成: {task_desc}")
+
+                        # 检查失败策略
+                        if not result.success and self.stop_on_failure:
+                            task_desc = item.description or name
+                            logger.warning(f"任务失败,停止执行: {task_desc}")
+                            stopped = True
+                            # 不再处理剩余任务
+                            break
+
+                    except (OSError, RuntimeError, ValueError) as e:
+                        task_desc = item.description or name
+                        logger.error(f"DAG任务异常: {task_desc}, 错误:{e}")
+                        results[name] = RunResult(
+                            returncode=-1,
+                            stdout="",
+                            stderr=str(e),
+                            execution_time=0.0,
+                            mode="failed",
+                            cmd=item.cmd if item.cmd else [],
+                        )
+                        stopped = self.stop_on_failure
+
+                    # 更新依赖(除非已停止)
+                    if not stopped:
+                        self._update_dependents(name, graph, in_degree, ready_queue)
+
+        return list(results.values())
+
+    @classmethod
+    def from_dict(cls, config: dict[str, Any]) -> TaskGroup:
+        """从字典配置创建任务组.
+
+        Parameters
+        ----------
+        config : dict[str, Any]
+            任务组配置字典
+
+        Returns
+        -------
+        TaskGroup
+            任务组实例
+        """
+        name = config.get("name", "unnamed-group")
+        description = config.get("description", "")
+        enabled = config.get("enabled", True)
+        parallel = config.get("parallel", False)
+        max_workers = config.get("max_workers")
+        stop_on_failure = config.get("stop_on_failure", False)
+
+        commands = []
+        for cmd_config in config.get("commands", []):
+            if isinstance(cmd_config, dict):
+                task_item = Task(
+                    name=cmd_config.get("name", ""),
+                    depends_on=cmd_config.get("depends_on", []),
+                    cmd=cmd_config.get("cmd"),
+                    func=cmd_config.get("func"),
+                    condition=cmd_config.get("condition"),
+                    description=cmd_config.get("description", ""),
+                    parallel=cmd_config.get("parallel", False),
+                )
+                commands.append(task_item)
+            elif isinstance(cmd_config, list):
+                # 直接是命令列表
+                commands.append(Task(cmd=cmd_config))
+            elif callable(cmd_config):
+                # 直接是可调用对象
+                commands.append(Task(func=cmd_config))
+
+        return cls(
+            name=name,
+            commands=commands,
+            description=description,
+            enabled=enabled,
+            parallel=parallel,
+            max_workers=max_workers,
+            stop_on_failure=stop_on_failure,
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """转换为字典配置."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "enabled": self.enabled,
+            "parallel": self.parallel,
+            "max_workers": self.max_workers,
+            "stop_on_failure": self.stop_on_failure,
+            "commands": [
+                {
+                    "name": item.name,
+                    "cmd": item.cmd,
+                    "depends_on": item.depends_on,
+                    "description": item.description,
+                    "parallel": item.parallel,
+                    # 注意: func 和 condition 无法序列化
+                }
+                for item in self.commands
+            ],
+        }