78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
# Copyright: Ankitects Pty Ltd and contributors
|
|
# License: GNU AGPL, version 3 or later; http://www.gnu.org/licenses/agpl.html
|
|
|
|
"""
|
|
Helper for running tasks on background threads.
|
|
"""
|
|
|
|
from concurrent.futures import Future
|
|
from concurrent.futures.thread import ThreadPoolExecutor
|
|
from dataclasses import dataclass
|
|
from threading import Lock
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
|
|
from PyQt5.QtCore import QObject, pyqtSignal
|
|
|
|
|
|
@dataclass
|
|
class PendingDoneCallback:
|
|
callback: Callable[[Any], None]
|
|
future: Future
|
|
|
|
|
|
class TaskManager(QObject):
|
|
_results_available = pyqtSignal()
|
|
|
|
def __init__(self):
|
|
QObject.__init__(self)
|
|
|
|
self._executor = ThreadPoolExecutor()
|
|
|
|
self._pending_callbacks: List[PendingDoneCallback] = []
|
|
|
|
self._results_lock = Lock()
|
|
self._results_available.connect(self._drain_results) # type: ignore
|
|
|
|
def run(
|
|
self,
|
|
task: Callable,
|
|
on_done: Optional[Callable],
|
|
args: Optional[Dict[str, Any]] = None,
|
|
) -> Future:
|
|
"""Run task on a background thread, calling on_done on the main thread if provided.
|
|
|
|
Args if provided will be passed on as keyword arguments to the task callable."""
|
|
if args is None:
|
|
args = {}
|
|
|
|
fut = self._executor.submit(task, **args)
|
|
|
|
if on_done is not None:
|
|
|
|
def done_closure(completed_future: Future) -> None:
|
|
self._handle_done_callback(completed_future, on_done)
|
|
|
|
fut.add_done_callback(done_closure)
|
|
|
|
return fut
|
|
|
|
def _handle_done_callback(self, future: Future, callback: Callable) -> None:
|
|
"""When future completes, schedule its callback to run on the main thread."""
|
|
# add result to the queue
|
|
with self._results_lock:
|
|
self._pending_callbacks.append(
|
|
PendingDoneCallback(callback=callback, future=future)
|
|
)
|
|
|
|
# and tell the main thread to flush the queue
|
|
self._results_available.emit() # type: ignore
|
|
|
|
def _drain_results(self):
|
|
"""Fires pending callbacks in the main thread."""
|
|
with self._results_lock:
|
|
results = self._pending_callbacks
|
|
self._pending_callbacks = []
|
|
|
|
for result in results:
|
|
result.callback(result.future)
|