bitool 0.1.2__py3-none-any.whl

bitool/utils/profiler.py

@@ -0,0 +1,369 @@
+"""性能分析工具模块.
+
+提供基于装饰器的性能分析功能,追踪函数运行时间和内存使用。
+
+示例:
+    >>> from bitool import profile
+    >>> @profile
+    ... def my_function():
+    ...     pass
+    >>> my_function()
+
+性能测试应用:
+    通过环境变量 BITOOL_PROFILE=1 启用性能分析,避免影响发布脚本性能。
+    在性能测试脚本中使用,而不是在生产脚本中使用。
+"""
+
+from __future__ import annotations
+
+import contextlib
+import functools
+import threading
+import time
+import tracemalloc
+from dataclasses import dataclass, field
+from functools import cached_property
+
+from rich.console import Console
+from rich.table import Table
+
+from ..core import env_check_enabled, logger
+from ..types import FunctionType, ParamType, ReturnType
+
+# 延迟初始化 tracemalloc,避免影响正常性能
+_tracemalloc_started = False
+
+
+# 常量定义
+_MS_PER_SEC = 1000  # 每秒毫秒数
+_μS_PER_MS = 1000  # 每毫秒微秒数
+_MEMORY_THRESHOLD_MB = 0.01  # 内存变化阈值(MB)
+_MAX_FUNC_NAME_LENGTH = 28  # 函数名最大长度
+_FUNC_NAME_TRUNCATE_LENGTH = 25  # 函数名截断长度
+
+# 显示配置常量
+_CONSOLE_WIDTH = 120  # 控制台宽度
+_TABLE_TITLE_STYLE = "bold magenta"  # 表格标题样式
+_TABLE_BORDER_STYLE = "blue"  # 表格边框样式
+_TABLE_HEADER_STYLE = "bold cyan"  # 表头样式
+_COL_FUNC_NAME_STYLE = "green"  # 函数名列样式
+_COL_DURATION_STYLE = "yellow"  # 耗时列样式
+_COL_PERCENTAGE_STYLE = "blue"  # 占比列样式
+_COL_CALL_COUNT_STYLE = "magenta"  # 调用次数列样式
+_COL_MEMORY_STYLE = "red"  # 内存列样式
+
+
+_USE_RUST_PROFILER = False
+
+
+__all__ = ["profile"]
+
+
+@dataclass
+class _PyProfileRecord:
+    """Python 实现的性能记录(用于 Rust 不可用时的降级)."""
+
+    name: str
+    duration_secs: float
+    memory_before_mb: float = 0.0
+    memory_after_mb: float = 0.0
+    memory_delta_mb: float = 0.0
+    timestamp: str = ""
+
+
+class _PythonPerformanceProfiler:
+    """Python 实现的性能分析器(用于 Rust 不可用时的降级)."""
+
+    def __init__(self) -> None:
+        self._records: list[_PyProfileRecord] = []
+        self._active_timers: dict[str, float] = {}
+        self._call_counts: dict[str, int] = {}
+        self._ensure_tracemalloc_started()
+
+    def _ensure_tracemalloc_started(self) -> None:
+        """确保 tracemalloc 已启动(延迟初始化)."""
+        global _tracemalloc_started
+        if not _tracemalloc_started:
+            try:
+                tracemalloc.start()
+                _tracemalloc_started = True
+                logger.debug("已启动 tracemalloc 内存追踪")
+            except RuntimeError as e:
+                logger.warning(f"启动 tracemalloc 失败: {e}")
+
+    def start_timer(self, name: str) -> None:
+        """开始计时."""
+        self._active_timers[name] = time.perf_counter()
+        # 记录当前内存使用
+        try:
+            current, _peak = tracemalloc.get_traced_memory()
+            self._active_timers[f"{name}_memory"] = current / (1024 * 1024)
+        except RuntimeError as e:
+            logger.debug(f"获取内存信息失败: {e}")
+
+    def stop_timer(self, name: str) -> _PyProfileRecord | None:
+        """停止计时并记录性能数据."""
+        start_time = self._active_timers.pop(name, None)
+        if start_time is None:
+            logger.warning(f"未找到计时器: {name}")
+            return None
+
+        duration = time.perf_counter() - start_time
+
+        # 获取内存信息
+        memory_before = self._active_timers.pop(f"{name}_memory", 0.0)
+        memory_after = 0.0
+        memory_delta = 0.0
+        try:
+            current, _peak = tracemalloc.get_traced_memory()
+            memory_after = current / (1024 * 1024)
+            memory_delta = memory_after - memory_before
+        except RuntimeError as e:
+            logger.debug(f"获取内存信息失败: {e}")
+
+        record = _PyProfileRecord(
+            name=name,
+            duration_secs=duration,
+            memory_before_mb=memory_before,
+            memory_after_mb=memory_after,
+            memory_delta_mb=memory_delta,
+            timestamp=time.strftime("%Y-%m-%d %H:%M:%S"),
+        )
+        self._records.append(record)
+        self._call_counts[name] = self._call_counts.get(name, 0) + 1
+        return record
+
+    def get_records(self) -> list[_PyProfileRecord]:
+        """获取所有性能记录."""
+        return self._records.copy()
+
+    def get_call_count(self, name: str) -> int:
+        """获取调用次数."""
+        return self._call_counts.get(name, 0)
+
+    def clear(self) -> None:
+        """清除所有记录."""
+        self._records.clear()
+        self._active_timers.clear()
+        self._call_counts.clear()
+
+
+@dataclass
+class _FuncStats:
+    """函数性能统计数据."""
+
+    total_time: float = 0.0  # 总耗时(包含子调用)
+    self_time: float = 0.0  # 自身耗时(不包含子调用)
+    call_count: int = 0
+    memory_delta: float = 0.0
+
+
+@dataclass
+class _PerformanceProfilerWrapper:
+    """性能分析器包装器."""
+
+    _call_depth: dict[int, int] = field(default_factory=dict)
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+    _child_time: dict[int, float] = field(default_factory=dict)
+    _self_times: dict[str, float] = field(default_factory=dict)
+
+    def profile(self, func: FunctionType) -> FunctionType:
+        """性能分析装饰器."""
+        # 动态检查环境变量,支持运行时启用
+        if not env_check_enabled("BITOOL_PROFILE"):
+            return func
+
+        func_name: str = getattr(func, "__name__", "Unknown")
+
+        @functools.wraps(func)
+        def wrapper(*args: ParamType.args, **kwargs: ParamType.kwargs) -> ReturnType:
+            thread_id: int = threading.current_thread().ident or 0
+
+            with self._lock:
+                self._call_depth.setdefault(thread_id, 0)
+                self._call_depth[thread_id] += 1
+                is_top_level = self._call_depth[thread_id] == 1
+                # 记录进入函数前的子调用耗时(用于计算自身耗时)
+                self._child_time.setdefault(thread_id, 0.0)
+                entry_child_time = self._child_time[thread_id]
+
+            self.profiler.start_timer(func_name)
+
+            try:
+                result = func(*args, **kwargs)
+            except BaseException:
+                # 异常情况下也要停止计时器, 避免计时器泄漏
+                with contextlib.suppress(BaseException):
+                    self.profiler.stop_timer(func_name)
+                raise
+            else:
+                # 正常执行完成,停止计时器
+                record = self.profiler.stop_timer(func_name)
+
+                # 计算自身耗时
+                if record is not None:
+                    # 退出时的子调用耗时 - 进入时的子调用耗时 = 当前函数的子调用耗时
+                    with self._lock:
+                        exit_child_time = self._child_time[thread_id]
+                        children_time = exit_child_time - entry_child_time
+                        self_time = record.duration_secs - children_time
+
+                    # 保存自身耗时到字典中
+                    self._self_times[func_name] = (
+                        self._self_times.get(func_name, 0.0) + self_time
+                    )
+
+                    # 将当前函数的总耗时累加到父函数的子调用耗时中
+                    with self._lock:
+                        if self._call_depth[thread_id] > 1:  # 如果有父函数
+                            self._child_time[thread_id] += record.duration_secs
+
+                return result
+            finally:
+                with self._lock:
+                    self._call_depth[thread_id] -= 1
+                    if is_top_level and self._call_depth[thread_id] == 0:
+                        self._show_summary()
+
+        return wrapper  # type: ignore[return-value]
+
+    def _show_summary(self) -> None:
+        """显示性能分析摘要."""
+        if not self.records:
+            logger.info("无性能数据")
+            return
+
+        # 使用最后一条记录(顶层函数)的耗时作为基准
+        baseline_time = self.records[-1].duration_secs
+
+        # 创建 Rich 表格(使用缓存的 Console 实例)
+        console = Console(width=_CONSOLE_WIDTH)
+        table = Table(
+            title="性能分析报告",
+            title_style=_TABLE_TITLE_STYLE,
+            border_style=_TABLE_BORDER_STYLE,
+            show_header=True,
+            header_style=_TABLE_HEADER_STYLE,
+        )
+
+        # 添加列
+        table.add_column("函数名", style=_COL_FUNC_NAME_STYLE, max_width=30)
+        table.add_column("总耗时", justify="right", style=_COL_DURATION_STYLE)
+        table.add_column("自身耗时", justify="right", style=_COL_DURATION_STYLE)
+        table.add_column("占比", justify="right", style=_COL_PERCENTAGE_STYLE)
+        table.add_column("调用次数", justify="right", style=_COL_CALL_COUNT_STYLE)
+        table.add_column("内存变化", justify="right", style=_COL_MEMORY_STYLE)
+
+        # 添加数据行
+        for name, stats in self.sorted_funcs:
+            # 格式化总耗时
+            total_time_str = self._format_duration(stats.total_time)
+            # 格式化自身耗时
+            self_time_str = self._format_duration(stats.self_time)
+
+            # 格式化占比(相对于顶层函数耗时,使用自身耗时)
+            percentage = (
+                stats.self_time / baseline_time * 100 if baseline_time > 0 else 0
+            )
+            percentage_str = f"{percentage:.1f}%"
+
+            # 格式化调用次数
+            call_str = f"{stats.call_count}x"
+
+            # 格式化内存变化
+            mem_str = self._format_memory(stats.memory_delta)
+
+            # 截断过长的函数名
+            display_name = self._truncate_func_name(name)
+
+            table.add_row(
+                display_name,
+                total_time_str,
+                self_time_str,
+                percentage_str,
+                call_str,
+                mem_str,
+            )
+
+        # 添加总计行
+        total_str = self._format_duration(baseline_time)
+
+        # 计算总调用次数
+        total_calls = sum(stats.call_count for _, stats in self.func_stats.items())
+
+        table.add_row(
+            "总计", total_str, "", "100.0%", f"{total_calls}次", "", style="bold"
+        )
+
+        # 输出表格
+        console.print(table)
+
+    @cached_property
+    def profiler(self) -> _PythonPerformanceProfiler:
+        """获取性能分析器实例."""
+        return _PythonPerformanceProfiler()
+
+    @cached_property
+    def records(self) -> list[_PyProfileRecord]:
+        """获取所有性能记录."""
+        return self.profiler.get_records()
+
+    @cached_property
+    def total_time(self) -> float:
+        """获取总耗时(秒)."""
+        return sum(record.duration_secs for record in self.records)
+
+    @cached_property
+    def func_stats(self) -> dict[str, _FuncStats]:
+        """获取所有函数的性能统计数据."""
+        from collections import defaultdict
+
+        func_stats: dict[str, _FuncStats] = defaultdict(_FuncStats)
+        for record in self.records:
+            name = record.name
+            func_stats[name].total_time += record.duration_secs
+            func_stats[name].call_count += 1
+            func_stats[name].memory_delta += record.memory_delta_mb
+            # 使用字典中的自身耗时
+            func_stats[name].self_time = self._self_times.get(name, 0.0)
+        return dict(func_stats)
+
+    @cached_property
+    def sorted_funcs(self) -> list[tuple[str, _FuncStats]]:
+        """获取按总耗时排序的函数性能统计数据."""
+        return sorted(
+            self.func_stats.items(), key=lambda x: x[1].total_time, reverse=True
+        )
+
+    def _format_duration(self, seconds: float) -> str:
+        """格式化时间显示."""
+        total_ms = seconds * _MS_PER_SEC
+        if total_ms >= _MS_PER_SEC:
+            return f"{total_ms / _MS_PER_SEC:.2f}s"
+        if total_ms >= 1:
+            return f"{total_ms:.0f}ms"
+        return f"{total_ms * _μS_PER_MS:.0f}μs"
+
+    def _format_memory(self, delta_mb: float) -> str:
+        """格式化内存变化显示."""
+        if abs(delta_mb) > _MEMORY_THRESHOLD_MB:
+            return f"{delta_mb:+.2f}MB"
+        return "~0MB"
+
+    def _truncate_func_name(self, name: str) -> str:
+        """截断过长的函数名."""
+        if len(name) > _MAX_FUNC_NAME_LENGTH:
+            return name[:_FUNC_NAME_TRUNCATE_LENGTH] + "..."
+        return name
+
+    def clear(self) -> None:
+        """清除所有性能记录."""
+        self.profiler.clear()
+        with self._lock:
+            self._call_depth.clear()
+            self._child_time.clear()
+            self._self_times.clear()
+
+
+_profiler_wrapper = _PerformanceProfilerWrapper()
+profile = _profiler_wrapper.profile

bitool/utils/task.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from functools import cached_property
+from typing import Any, Callable
+
+from .executor import RunResult
+
+
+@dataclass
+class Task:
+    """单个任务项.
+
+    Attributes
+    ----------
+    name : str
+        任务唯一标识,用于依赖引用
+    cmd : list[str], optional
+        命令及其参数列表
+    func : Callable[[], RunResult], optional
+        可执行的函数
+    condition : Callable[[], bool], optional
+        执行条件函数,返回True才执行
+    description : str
+        任务描述信息
+    parallel : bool
+        是否允许并行执行,默认 False (顺序执行)
+    depends_on : list[str]
+        依赖的任务名称列表,用于DAG调度
+    """
+
+    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] = field(default_factory=list)
+
+    def __post_init__(self) -> None:
+        """验证任务项配置."""
+        if self.cmd is None and self.func is None:
+            msg = "Task 必须提供 cmd 或 func 参数"
+            raise ValueError(msg)
+
+        # 如果声明了依赖,必须有名称
+        if self.depends_on and not self.name:
+            msg = "Task 如果声明了 depends_on,必须提供 name 参数"
+            raise ValueError(msg)
+
+    @cached_property
+    def is_executable(self) -> bool:
+        """检查是否应该执行此任务."""
+        if self.condition is None:
+            return True
+
+        return self.condition()
+
+    @cached_property
+    def executable_cmd(self) -> list[str] | Callable[[], RunResult] | None:
+        """转换为可执行对象."""
+        if self.cmd is not None:
+            return self.cmd
+
+        return self.func
+
+    @classmethod
+    def from_dict(cls, config: dict[str, Any]) -> Task:
+        """从字典配置创建任务.
+
+        Parameters
+        ----------
+        config : dict[str, Any]
+            任务配置字典
+
+        Returns
+        -------
+        Task
+            任务实例
+        """
+        return cls(
+            name=config.get("name", ""),
+            cmd=config.get("cmd"),
+            func=config.get("func"),
+            condition=config.get("condition"),
+            description=config.get("description", ""),
+            parallel=config.get("parallel", False),
+            depends_on=config.get("depends_on", []),
+        )
+
+    @classmethod
+    def from_callable(
+        cls,
+        func: Callable[[], RunResult],
+        description: str = "",
+    ) -> Task:
+        """从可调用对象创建任务.
+
+        Parameters
+        ----------
+        func : Callable[[], RunResult]
+            可执行函数
+        description : str
+            任务描述
+
+        Returns
+        -------
+        Task
+            任务实例
+        """
+        return cls(func=func, description=description)
+
+    @classmethod
+    def from_list(
+        cls,
+        cmd: list[str],
+        description: str = "",
+    ) -> Task:
+        """从命令列表创建任务.
+
+        Parameters
+        ----------
+        cmd : list[str]
+            命令及其参数列表
+        description : str
+            任务描述
+
+        Returns
+        -------
+        Task
+            任务实例
+        """
+        return cls(cmd=cmd, description=description)