zsglpt-pc/utils/worker.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
后台线程管理模块
使用QThread实现非阻塞的后台任务
"""

from typing import Callable, Any, Optional
from PyQt6.QtCore import QObject, QThread, pyqtSignal


class WorkerSignals(QObject):
    """工作线程信号类"""
    # 进度信号：(百分比, 消息)
    progress = pyqtSignal(int, str)
    # 日志信号：日志消息
    log = pyqtSignal(str)
    # 完成信号：(成功/失败, 结果/错误消息)
    finished = pyqtSignal(bool, str)
    # 截图完成信号：截图文件路径
    screenshot_ready = pyqtSignal(str)
    # 通用结果信号：任意结果对象
    result = pyqtSignal(object)
    # 错误信号
    error = pyqtSignal(str)


class Worker(QThread):
    """
    通用后台工作线程

    用法:
        worker = Worker(some_function, arg1, arg2, kwarg1=value1)
        worker.signals.finished.connect(on_finished)
        worker.signals.progress.connect(on_progress)
        worker.start()
    """

    def __init__(self, fn: Callable, *args, **kwargs):
        super().__init__()
        self.fn = fn
        self.args = args
        self.kwargs = kwargs
        self.signals = WorkerSignals()
        self._is_stopped = False

    def run(self):
        """执行后台任务"""
        try:
            # 把信号传递给任务函数，方便回调
            self.kwargs['_signals'] = self.signals
            self.kwargs['_should_stop'] = self.should_stop

            result = self.fn(*self.args, **self.kwargs)

            if not self._is_stopped:
                self.signals.result.emit(result)
                self.signals.finished.emit(True, "完成")
        except Exception as e:
            if not self._is_stopped:
                error_msg = str(e)
                self.signals.error.emit(error_msg)
                self.signals.finished.emit(False, error_msg)

    def stop(self):
        """停止线程"""
        self._is_stopped = True

    def should_stop(self) -> bool:
        """检查是否应该停止"""
        return self._is_stopped


class TaskRunner:
    """
    任务运行器
    管理多个后台任务，确保不会同时运行太多任务
    """

    def __init__(self, max_workers: int = 1):
        self.max_workers = max_workers
        self._workers: list[Worker] = []
        self._queue: list[tuple] = []  # (fn, args, kwargs, callbacks)

    def submit(self, fn: Callable, *args,
               on_progress: Optional[Callable] = None,
               on_log: Optional[Callable] = None,
               on_finished: Optional[Callable] = None,
               on_result: Optional[Callable] = None,
               on_error: Optional[Callable] = None,
               **kwargs) -> Optional[Worker]:
        """
        提交任务

        Args:
            fn: 要执行的函数
            *args: 位置参数
            on_progress: 进度回调 (percent, message)
            on_log: 日志回调 (message)
            on_finished: 完成回调 (success, message)
            on_result: 结果回调 (result)
            on_error: 错误回调 (error_message)
            **kwargs: 关键字参数

        Returns:
            Worker对象，如果队列满了返回None
        """
        callbacks = {
            'on_progress': on_progress,
            'on_log': on_log,
            'on_finished': on_finished,
            'on_result': on_result,
            'on_error': on_error,
        }

        # 清理已完成的worker
        self._workers = [w for w in self._workers if w.isRunning()]

        if len(self._workers) >= self.max_workers:
            # 加入队列等待
            self._queue.append((fn, args, kwargs, callbacks))
            return None

        return self._start_worker(fn, args, kwargs, callbacks)

    def _start_worker(self, fn: Callable, args: tuple, kwargs: dict,
                      callbacks: dict) -> Worker:
        """启动一个worker"""
        worker = Worker(fn, *args, **kwargs)

        # 连接信号
        if callbacks.get('on_progress'):
            worker.signals.progress.connect(callbacks['on_progress'])
        if callbacks.get('on_log'):
            worker.signals.log.connect(callbacks['on_log'])
        if callbacks.get('on_result'):
            worker.signals.result.connect(callbacks['on_result'])
        if callbacks.get('on_error'):
            worker.signals.error.connect(callbacks['on_error'])

        # 完成时的处理
        def on_worker_finished(success, message):
            if callbacks.get('on_finished'):
                callbacks['on_finished'](success, message)
            # 处理队列中的下一个任务
            self._process_queue()

        worker.signals.finished.connect(on_worker_finished)

        self._workers.append(worker)
        worker.start()
        return worker

    def _process_queue(self):
        """处理队列中的任务"""
        # 清理已完成的worker
        self._workers = [w for w in self._workers if w.isRunning()]

        if self._queue and len(self._workers) < self.max_workers:
            fn, args, kwargs, callbacks = self._queue.pop(0)
            self._start_worker(fn, args, kwargs, callbacks)

    def stop_all(self):
        """停止所有任务"""
        for worker in self._workers:
            worker.stop()
        self._queue.clear()

    def is_running(self) -> bool:
        """检查是否有任务在运行"""
        return any(w.isRunning() for w in self._workers)

    @property
    def running_count(self) -> int:
        """运行中的任务数"""
        return sum(1 for w in self._workers if w.isRunning())

    @property
    def queue_size(self) -> int:
        """队列中等待的任务数"""
        return len(self._queue)


# 全局任务运行器
_task_runner: Optional[TaskRunner] = None


def get_task_runner() -> TaskRunner:
    """获取全局任务运行器"""
    global _task_runner
    if _task_runner is None:
        _task_runner = TaskRunner(max_workers=1)
    return _task_runner