bitool 0.1.2__py3-none-any.whl

bitool/cmd/run.py

@@ -0,0 +1,300 @@
+from __future__ import annotations
+
+import shlex
+import subprocess
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, List, cast
+
+from bitool.cmd import BaseCommand
+from bitool.core import logger
+from bitool.utils.executor import RunResult, execute
+from bitool.utils.task import Task
+from bitool.utils.task_group import TaskGroup
+
+
+def _parse_shell_command(command: str) -> tuple[list[str], str | None, bool, bool]:
+    """安全地解析包含重定向的 shell 命令。
+
+    手动解析常见的 shell 操作符（>, >>, |），避免使用 shell=True。
+
+    Args:
+        command: shell 命令字符串。
+
+    Returns:
+        元组 (cmd_args, stdout_redirect, append_mode, has_pipe)
+        - cmd_args: 命令及其参数列表
+        - stdout_redirect: 标准输出重定向文件路径（如果有）
+        - append_mode: 是否为追加模式（>>）
+        - has_pipe: 是否包含管道（暂不支持）
+
+    Raises:
+        ValueError: 当命令包含不支持的操作符时。
+    """
+    # 检查管道操作符（暂不支持）
+    if "|" in command:
+        msg = "不支持管道操作符, 请使用多个命令分别执行"
+        raise ValueError(msg)
+
+    # 解析重定向操作符
+    stdout_redirect: str | None = None
+    append_mode = False
+
+    # 处理 >> (追加)
+    if ">>" in command:
+        parts_redirect: list[str] = command.split(">>", 1)
+        cmd_main: str = parts_redirect[0].strip()
+        stdout_redirect = parts_redirect[1].strip()
+        append_mode = True
+    # 处理 > (覆盖)
+    elif ">" in command:
+        parts_redirect = command.split(">", 1)
+        cmd_main = parts_redirect[0].strip()
+        stdout_redirect = parts_redirect[1].strip()
+    else:
+        cmd_main = command
+
+    # 使用 shlex 安全地分割命令
+    cmd_args: list[str] = shlex.split(cmd_main)
+
+    return cmd_args, stdout_redirect, append_mode, False
+
+
+__all__ = [
+    "RunCommand",
+    "RunDAGCommands",
+    "RunParallelCommands",
+    "RunResult",
+    "RunSequentialCommands",
+    "RunShellCommand",
+]
+
+
+@dataclass
+class BaseRunCommand(BaseCommand):
+    cmd: list[str] | Callable[[], None] = field(default_factory=list)
+    cwd: Path | None = None
+    env: dict[str, str] | None = None
+    timeout: int | None = None
+    check: bool = False
+
+
+@dataclass
+class RunCommand(BaseRunCommand):
+    """执行单个命令"""
+
+    name = "run"
+    description = "执行单个命令"
+
+    def run(self) -> bool:
+        """执行单个命令"""
+        if isinstance(self.cmd, list):
+            cmd_list = cast(List[str], self.cmd)
+            cmd_str = " ".join(str(arg) for arg in cmd_list)
+            try:
+                execute(
+                    cmd_list,
+                    cwd=self.cwd,
+                    env=self.env,
+                    timeout=self.timeout,
+                    check=self.check,
+                )
+                # 使用列表推导确保类型安全
+                logger.info(f"命令执行成功: `{cmd_str}`")
+            except Exception as e:  # noqa: BLE001
+                logger.error(f"命令执行失败: `{cmd_str}`, 错误: {e}")
+                return False
+        else:
+            # 命令是 Callable[[], None] 类型，直接调用
+            self.cmd()
+        return True
+
+
+@dataclass
+class RunShellCommand(BaseRunCommand):
+    """执行shell命令字符串"""
+
+    name = "run_shell"
+    description = "执行shell命令字符串"
+    shell_cmd: str = ""
+
+    def run(self) -> bool:
+        """执行shell命令字符串（安全版本，支持重定向）。"""
+        try:
+            # 安全解析命令（避免使用 shell=True）
+            cmd_args, stdout_redirect, append_mode, _ = _parse_shell_command(
+                self.shell_cmd
+            )
+
+            if not cmd_args:
+                logger.error(f"Shell命令为空: `{self.shell_cmd}`")
+                return False
+
+            # 执行命令并处理重定向
+            if stdout_redirect:
+                # 确保输出目录存在
+                redirect_path = Path(stdout_redirect)
+                if redirect_path.parent != Path("."):
+                    redirect_path.parent.mkdir(parents=True, exist_ok=True)
+
+                # 以追加或覆盖模式打开文件
+                mode = "a" if append_mode else "w"
+                with redirect_path.open(mode, encoding="utf-8") as outfile:
+                    result = subprocess.run(  # noqa: S603
+                        cmd_args,
+                        cwd=self.cwd,
+                        timeout=self.timeout,
+                        stdout=outfile,
+                        stderr=subprocess.PIPE,
+                        text=True,
+                        check=False,
+                    )
+
+                    if result.returncode == 0:
+                        logger.info(f"Shell命令执行成功: `{self.shell_cmd}`")
+                        return True
+                    else:
+                        error_output = result.stderr.strip()
+                        logger.error(
+                            f"Shell命令执行失败: `{self.shell_cmd}`, 返回码: {result.returncode}, 错误: {error_output}"
+                        )
+                        return False
+            else:
+                # 无重定向，正常执行
+                result = subprocess.run(  # noqa: S603
+                    cmd_args,
+                    cwd=self.cwd,
+                    timeout=self.timeout,
+                    capture_output=True,
+                    text=True,
+                    check=False,
+                )
+
+                if result.returncode == 0:
+                    logger.info(f"Shell命令执行成功: `{self.shell_cmd}`")
+                    if result.stdout.strip():
+                        for line in result.stdout.strip().split("\n"):
+                            if line.strip():
+                                logger.info(f"  输出: `{line.strip()}`")
+                    return True
+                else:
+                    error_output = result.stderr.strip()
+                    logger.error(
+                        f"Shell命令执行失败: `{self.shell_cmd}`, 返回码: {result.returncode}, 错误: {error_output}"
+                    )
+                    return False
+
+        except subprocess.TimeoutExpired:
+            logger.error(f"Shell命令执行超时: `{self.shell_cmd}`")
+            return False
+        except ValueError as e:
+            logger.error(f"Shell命令解析失败: `{self.shell_cmd}`, 错误: {e}")
+            return False
+        except Exception as e:  # noqa: BLE001
+            logger.error(f"Shell命令执行异常: `{self.shell_cmd}`, 错误: {e}")
+            return False
+
+
+@dataclass
+class BaseRunMultipleCommands(BaseCommand):
+    commands: list[list[str]] = field(default_factory=list)
+    stop_on_failure: bool = False
+    max_workers: int | None = None
+    results: list[RunResult] = field(default_factory=list)
+
+
+@dataclass
+class RunSequentialCommands(BaseRunMultipleCommands):
+    """顺序执行多个命令"""
+
+    name = "run_sequential"
+    description = "顺序执行多个命令"
+
+    def run(self) -> bool:
+        """顺序执行多个命令"""
+        if not self.commands:
+            logger.info("命令列表为空,直接返回")
+            return True
+
+        tasks = [Task(cmd=cmd) for cmd in self.commands]
+        group = TaskGroup(
+            name="sequential-execution",
+            commands=tasks,
+            parallel=False,
+            stop_on_failure=self.stop_on_failure,
+        )
+
+        logger.info(f"开始顺序执行 {len(self.commands)} 个命令")
+        self.results = group.execute()
+        return all(r.returncode == 0 for r in self.results)
+
+
+@dataclass
+class RunParallelCommands(BaseRunMultipleCommands):
+    """并行执行多个命令"""
+
+    name = "run_parallel"
+    description = "并行执行多个命令"
+
+    def run(self) -> bool:
+        """并行执行多个命令"""
+        if not self.commands:
+            logger.info("命令列表为空,直接返回")
+            return True
+
+        tasks = [Task(cmd=cmd, parallel=True) for cmd in self.commands]
+        group = TaskGroup(
+            name="parallel-execution",
+            commands=tasks,
+            parallel=True,
+            max_workers=self.max_workers,
+            stop_on_failure=False,
+        )
+
+        logger.info(f"开始并行执行 {len(self.commands)} 个命令")
+        self.results = group.execute()
+        return all(r.returncode == 0 for r in self.results)
+
+
+@dataclass
+class RunDAGCommands(BaseRunMultipleCommands):
+    """基于DAG依赖调度执行命令"""
+
+    name = "run_dag"
+    description = "基于DAG依赖调度执行命令"
+    tasks: list[dict[str, object]] = field(default_factory=list)
+
+    def run(self) -> bool:
+        """基于DAG依赖调度执行命令"""
+        if not self.tasks:
+            logger.info("任务列表为空,直接返回")
+            return True
+
+        task_objects = []
+        for task_config in self.tasks:
+            # 类型守卫: 从 dict 中提取具体类型
+            name = str(task_config.get("name", ""))
+            cmd = task_config.get("cmd")
+            func = task_config.get("func")
+            depends_on_val = task_config.get("depends_on", [])
+            description = str(task_config.get("description", ""))
+
+            task = Task(
+                name=name,
+                cmd=cmd if isinstance(cmd, list) else None,  # type: ignore
+                func=func if callable(func) else None,  # type: ignore
+                depends_on=depends_on_val if isinstance(depends_on_val, list) else [],  # type: ignore
+                description=description,
+            )
+            task_objects.append(task)
+
+        group = TaskGroup(
+            name="dag-execution",
+            commands=task_objects,
+            max_workers=self.max_workers,
+            stop_on_failure=self.stop_on_failure,
+        )
+
+        logger.info(f"开始DAG调度执行 {len(self.tasks)} 个任务")
+        self.results = group.execute()
+        return all(r.returncode == 0 for r in self.results)

bitool/cmd/toml.py

@@ -0,0 +1,237 @@
+"""TOML文件操作命令模块.
+
+提供TOML文件的读取、写入和版本号更新等功能.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+try:
+    import tomli
+    import tomli_w
+except ImportError:
+    try:
+        import tomli_w
+        import tomllib as tomli  # ty: ignore[unresolved-import]
+    except ImportError:
+        tomli = None  # ty: ignore
+        tomli_w = None  # ty: ignore
+
+from bitool.cmd._base import BaseCommand
+from bitool.core import logger
+
+
+@dataclass
+class ReadTomlCommand(BaseCommand):
+    """读取TOML文件命令.
+
+    Attributes:
+        filepath: TOML文件路径
+        data: 解析后的TOML数据
+    """
+
+    name: str = "read-toml"
+    description: str = "读取TOML文件"
+    filepath: Path = Path()
+    data: dict[str, Any] = field(default_factory=dict)
+
+    def run(self) -> bool:
+        """执行TOML文件读取.
+
+        Returns
+        -------
+        bool
+            读取成功返回True, 否则返回False
+        """
+        if not self.filepath.exists():
+            logger.warning(f"TOML文件不存在: {self.filepath}")
+            return False
+
+        try:
+            with self.filepath.open("rb") as f:
+                if tomli is not None:
+                    self.data = tomli.load(f)
+                else:
+                    logger.error("未安装tomli或tomllib库")
+                    return False
+        except Exception as e:  # noqa: BLE001
+            logger.error(f"读取TOML文件失败: {e}")
+            return False
+        else:
+            logger.info(f"成功读取TOML文件: {self.filepath}")
+            return True
+
+
+@dataclass
+class WriteTomlCommand(BaseCommand):
+    """写入TOML文件命令.
+
+    Attributes:
+        filepath: TOML文件路径
+        data: 要写入的TOML数据
+        overwrite: 是否覆盖已存在的文件
+        backup: 写入前是否备份原文件
+    """
+
+    name: str = "write-toml"
+    description: str = "写入TOML文件"
+    filepath: Path = Path()
+    data: dict[str, Any] = field(default_factory=dict)
+    overwrite: bool = True
+    backup: bool = False
+
+    def run(self) -> bool:
+        """执行TOML文件写入.
+
+        Returns
+        -------
+        bool
+            写入成功返回True, 否则返回False
+        """
+        if tomli_w is None:
+            logger.error("未安装tomli_w库")
+            return False
+
+        if self.filepath.exists() and not self.overwrite:
+            logger.warning(f"TOML文件已存在且不允许覆盖: {self.filepath}")
+            return False
+
+        try:
+            # 备份原文件
+            if self.backup and self.filepath.exists():
+                backup_path = self.filepath.with_suffix(self.filepath.suffix + ".bak")
+                backup_path.write_bytes(self.filepath.read_bytes())
+                logger.info(f"已备份TOML文件: {self.filepath} -> {backup_path}")
+
+            # 确保父目录存在
+            if not self.filepath.parent.exists():
+                self.filepath.parent.mkdir(parents=True, exist_ok=True)
+
+            # 写入TOML文件
+            with self.filepath.open("wb") as f:
+                tomli_w.dump(self.data, f)
+        except Exception as e:  # noqa: BLE001
+            logger.error(f"写入TOML文件失败: {e}")
+            return False
+        else:
+            logger.info(f"成功写入TOML文件: {self.filepath}")
+            return True
+
+
+@dataclass
+class UpdateTomlVersionCommand(BaseCommand):
+    """更新TOML文件版本号命令.
+
+    Attributes:
+        filepath: TOML文件路径
+        version: 新版本号字符串
+        version_path: 版本号在TOML中的路径, 如 ["project", "version"]
+    """
+
+    name: str = "update-toml-version"
+    description: str = "更新TOML文件中的版本号"
+    filepath: Path = Path()
+    version: str = ""
+    version_path: list[str] = field(default_factory=lambda: ["project", "version"])
+
+    def run(self) -> bool:
+        """执行TOML文件版本号更新.
+
+        Returns
+        -------
+        bool
+            更新成功返回True, 否则返回False
+        """
+        if tomli is None or tomli_w is None:
+            logger.error("未安装tomli或tomli_w库")
+            return False
+
+        if not self.filepath.exists():
+            logger.warning(f"TOML文件不存在: {self.filepath}")
+            return False
+
+        try:
+            # 读取原文件
+            with self.filepath.open("rb") as f:
+                data = tomli.load(f)
+
+            # 导航到版本号位置并更新
+            current = data
+            for key in self.version_path[:-1]:
+                if key not in current:
+                    current[key] = {}
+                current = current[key]
+
+            old_version = current.get(self.version_path[-1], "未知")
+            current[self.version_path[-1]] = self.version
+
+            # 写回文件
+            with self.filepath.open("wb") as f:
+                tomli_w.dump(data, f)
+
+            logger.info(
+                f"已更新版本号: {self.filepath} ({old_version} -> {self.version})"
+            )
+        except Exception as e:  # noqa: BLE001
+            logger.error(f"更新TOML版本号失败: {e}")
+            return False
+        else:
+            return True
+
+
+@dataclass
+class ReadTomlVersionCommand(BaseCommand):
+    """读取TOML文件版本号命令.
+
+    Attributes:
+        filepath: TOML文件路径
+        version: 读取到的版本号
+        version_path: 版本号在TOML中的路径
+    """
+
+    name: str = "read-toml-version"
+    description: str = "读取TOML文件中的版本号"
+    filepath: Path = Path()
+    version: str = ""
+    version_path: list[str] = field(default_factory=lambda: ["project", "version"])
+
+    def run(self) -> bool:
+        """执行TOML文件版本号读取.
+
+        Returns
+        -------
+        bool
+            读取成功返回True, 否则返回False
+        """
+        if tomli is None:
+            logger.error("未安装tomli或tomllib库")
+            return False
+
+        if not self.filepath.exists():
+            logger.warning(f"TOML文件不存在: {self.filepath}")
+            return False
+
+        try:
+            with self.filepath.open("rb") as f:
+                data = tomli.load(f)
+
+            # 导航到版本号位置
+            current = data
+            for key in self.version_path:
+                if key not in current:
+                    logger.warning(
+                        f"在 {self.filepath} 中未找到版本号路径: {'.'.join(self.version_path)}"
+                    )
+                    return False
+                current = current[key]
+
+            self.version = str(current)
+            logger.info(f"读取到版本号: {self.filepath} -> {self.version}")
+        except Exception as e:  # noqa: BLE001
+            logger.error(f"读取TOML版本号失败: {e}")
+            return False
+        else:
+            return True