File: C:/Users/fred/anaconda3/Lib/site-packages/anaconda_navigator/utils/workers.py
# -*- coding: utf-8 -*-
"""Integration with worker framework to make API calls asynchronous."""
from __future__ import annotations
__all__ = [
'TaskCanceledError', 'TaskFailedError', 'Call', 'Duration',
'TaskStatus', 'TaskResult', 'TaskWorker', 'TaskThread',
'TaskContext', 'TaskModifier', 'InsertArgument', 'Task', 'CancelContext', 'AddCancelContext',
]
import abc
import contextlib
import datetime
import enum
import typing
from qtpy import QtCore
from anaconda_navigator.utils.logs import logger
from . import singletons
ACTION = typing.Callable[[], typing.Any]
FUNCTION = typing.Callable[..., typing.Any]
class TaskError(Exception):
"""General issue running a task."""
class TaskCanceledError(TaskError):
"""Task was canceled during its workflow."""
def __str__(self) -> str:
"""Prepare string representation of the instance."""
return 'Task was cancelled during its execution'
class TaskFailedError(TaskError):
"""Task failed during its workflow."""
def __str__(self) -> str:
"""Prepare string representation of the instance."""
return 'Task failed during its execution'
class Call:
"""
Signature of a function call.
:param function: Function to call.
:param args: Positional arguments for the `function`.
:param kwargs: Keyword arguments for the `function`.
"""
__slots__ = ('__function', '__args', '__kwargs')
def __init__(
self,
function: FUNCTION,
args: typing.Optional[typing.Sequence[typing.Any]] = None,
kwargs: typing.Optional[typing.Mapping[str, typing.Any]] = None,
) -> None:
"""Initialize new :class:`~Call` instance."""
if args is None:
args = ()
if kwargs is None:
kwargs = {}
self.__function: typing.Final[FUNCTION] = function
self.__args: typing.Final[typing.Sequence[typing.Any]] = args
self.__kwargs: typing.Final[typing.Mapping[str, typing.Any]] = kwargs
@property
def function(self) -> FUNCTION: # noqa: D401
"""Function to call."""
return self.__function
@property
def args(self) -> typing.Sequence[typing.Any]: # noqa: D401
"""Positional arguments to call `function` with."""
return self.__args
@property
def kwargs(self) -> typing.Mapping[str, typing.Any]: # noqa: D401
"""Keyword arguments to call `function` with."""
return self.__kwargs
def __call__(self) -> typing.Any:
"""Execute stored function with both positional and keyword arguments."""
return self.__function(*self.__args, **self.__kwargs)
class Duration:
"""
Simple helper to track execution duration.
:param start: Optional start of the tracked time span.
If it is not provided - current timestamp will be used.
:param stop: Optional end of the tracked time span.
Might be omitted on initialization and then set with :meth:`~Duration.finish`.
"""
__slots__ = ('__start', '__stop')
def __init__(
self,
start: typing.Optional[datetime.datetime] = None,
stop: typing.Optional[datetime.datetime] = None,
) -> None:
"""Initialize new :class:`~Duration` instance."""
if start is None:
start = datetime.datetime.now()
self.__start: typing.Final[datetime.datetime] = start
self.__stop: typing.Optional[datetime.datetime] = stop
@property
def start(self) -> datetime.datetime: # noqa: D401
"""Start of the tracked time span."""
return self.__start
@property
def stop(self) -> typing.Optional[datetime.datetime]: # noqa: D401
"""Stop of the tracked time span."""
return self.__stop
@property
def total(self) -> typing.Optional[float]: # noqa: D401
"""Total number of seconds in this time span."""
if self.__stop is None:
return None
return (self.__stop - self.__start).total_seconds()
def finish(self, stop: typing.Optional[datetime.datetime] = None) -> None:
"""
Set the stop value for this duration.
Might be set only once (either with this method or on initialization).
:param stop: Value to set stop to.
If it is not provided - current timestamp will be used.
"""
if self.__stop is not None:
raise TypeError('Duration is already closed')
if stop is None:
stop = datetime.datetime.now()
self.__stop = stop
class TaskStatus(enum.Enum):
"""Possible states of the :class:`~Task` execution."""
SUCCEEDED = enum.auto()
CANCELED = enum.auto()
FAILED = enum.auto()
class TaskResult:
"""
Details of task execution.
Instances of this class are provided through the task signals.
:param call: Signature of the function that was called.
:param status: Resulted status of the execution.
:param result: Function output.
Must be an exact value for succeeded executions, and exception for failed executions.
:param duration: Optional information about task execution duration.
"""
__slots__ = ('__call', '__status', '__result', '__duration')
def __init__(
self,
call: Call,
status: TaskStatus,
result: typing.Any = None,
duration: typing.Optional[Duration] = None,
) -> None:
"""Initialize new :class:`~TaskResult` instance."""
if duration is None:
duration = Duration()
duration.finish()
if status == TaskStatus.CANCELED:
if result is None:
result = TaskCanceledError()
elif not isinstance(result, TaskCanceledError):
raise TypeError('result must be a TaskCanceledError instance if task is failed')
elif (status == TaskStatus.FAILED) and (not isinstance(result, BaseException)):
raise TypeError('result must be an exception if task is failed')
self.__call: typing.Final[Call] = call
self.__status: typing.Final[TaskStatus] = status
self.__result: typing.Final[typing.Any] = result
self.__duration: typing.Final[Duration] = duration
@property
def call(self) -> Call: # noqa: D401
"""Signature of the function that was called."""
return self.__call
@property
def duration(self) -> Duration: # noqa: D401
"""Task execution duration."""
return self.__duration
@property
def exception(self) -> typing.Optional[BaseException]: # noqa: D401
"""Exception, that was raised by the function."""
if self.__status == TaskStatus.SUCCEEDED:
return None
return self.__result
@property
def result(self) -> typing.Any: # noqa: D401
"""
Function output.
If function raised an error - this property will also raise the same error.
"""
if self.__status == TaskStatus.SUCCEEDED:
return self.__result
raise self.__result
@property
def status(self) -> TaskStatus: # noqa: D401
"""Resulted status of the execution."""
return self.__status
class TaskSignals(QtCore.QObject): # pylint: disable=too-few-public-methods
"""
Collection of signals for :class:`~TaskWorker`.
.. py:attribute:: sig_start
Task started its execution.
.. py:attribute:: sig_succeeded
Task finished its execution successfully.
.. py:attribute:: sig_canceled
Task was canceled during its execution.
.. py:attribute:: sig_failed
Task raised an error.
.. py:attribute:: sig_done
Task finished its execution.
This signal is called regardless of execution result.
"""
sig_start = QtCore.Signal()
sig_succeeded = QtCore.Signal(TaskResult)
sig_canceled = QtCore.Signal(TaskResult)
sig_failed = QtCore.Signal(TaskResult)
sig_done = QtCore.Signal(TaskResult)
class TaskWorker(QtCore.QRunnable):
"""
Executor for tasks.
:param call: Signature of the function to call with the worker.
:param cancel: Optional cancel function.
This function might be called to change `function` execution context, so it would be aware that it
should be canceled.
If it is not provided - `function` execution will still be finished gracefully, but result will
contain only information about cancelled function call, unless it will raise an exception.
"""
def __init__(self, call: Call, cancel: typing.Optional[ACTION] = None) -> None:
"""Initialize new :class:`~TaskWorker` instance."""
super().__init__()
self.__signals: typing.Final[TaskSignals] = TaskSignals()
self.__call: typing.Final[Call] = call
self.__cancel: typing.Final[typing.Optional[ACTION]] = cancel
self.__canceled: bool = False
self.__result: typing.Optional[TaskResult] = None
@property
def call(self) -> Call: # noqa: D401
"""Signature of the function that is executed in this task."""
return self.__call
@property
def signals(self) -> TaskSignals: # noqa: D401
"""Collection of signals to track task execution state."""
return self.__signals
@property
def result(self) -> typing.Optional[TaskResult]: # noqa: D401
"""Task execution result."""
return self.__result
def cancel(self) -> None:
"""Cancel execution of current task."""
if self.__result is not None:
raise TypeError('Already finished')
if self.__canceled:
raise TypeError('Already canceled')
self.__canceled = True
if self.__cancel is not None:
self.__cancel()
def run(self) -> None:
"""
Execute task.
This method is called internally by the Qt.
:meta private:
"""
MANAGER.instance._register(self) # pylint: disable=protected-access
self.signals.sig_start.emit()
status: TaskStatus
response: typing.Any
duration: typing.Final[Duration] = Duration()
try:
if self.__canceled:
raise TaskCanceledError()
response = self.__call()
status = TaskStatus.SUCCEEDED
except TaskCanceledError as exception:
response = exception
status = TaskStatus.CANCELED
except Exception as exception: # pylint: disable=broad-except
logger.debug('exception raised in worker', exc_info=True)
response = exception
status = TaskStatus.FAILED
finally:
duration.finish()
if self.__canceled:
response = TaskCanceledError()
status = TaskStatus.CANCELED
self.__result = TaskResult(
call=self.__call,
status=status,
result=response,
duration=duration,
)
{
TaskStatus.SUCCEEDED: self.signals.sig_succeeded,
TaskStatus.CANCELED: self.signals.sig_canceled,
TaskStatus.FAILED: self.signals.sig_failed,
}[status].emit(self.__result)
self.signals.sig_done.emit(self.__result)
def start(self) -> None:
"""Launch this task."""
THREAD_POOL.instance.start(self)
def thread(self) -> 'TaskThread':
"""Start a manageable thread with this task."""
result: TaskThread = TaskThread(parent=self)
result.start()
return result
class TaskThread(QtCore.QThread):
"""
Wrapper for a :class:`~TaskWorker`, that provides a manageable :class:`~qtpy.QtCore.QThread` interface.
Should be initialized from :class:`~TaskWorker` instance via :meth:`~TaskWorker.thread`.
"""
def __init__(self, parent: TaskWorker) -> None:
"""Initialize new :class:`~TaskThread` instance."""
super().__init__()
self.__parent: typing.Final[TaskWorker] = parent
@property
def call(self) -> Call: # noqa: D401
"""Signature of the function that is executed in this task."""
return self.__parent.call
@property
def result(self) -> typing.Optional[TaskResult]: # noqa: D401
"""Task execution result."""
return self.__parent.result
def cancel(self) -> None:
"""Cancel execution of current task."""
self.__parent.cancel()
def run(self) -> None:
"""
Execute the thread.
This method is called internally by the Qt.
:meta private:
"""
self.__parent.run()
class TaskContext:
"""
Context for task signature.
This is used to prepare contexts for both function :class:`~Call` and/or :class:`~TaskWorker`.
:param function: Function that should be called.
:param args: Initial collection of positional arguments to call function with.
:param kwargs: Initial collection of keyword arguments to call function with.
:param cancel: Initial value for cancel function, which may be used by :class:`~TaskWorker`.
Values of `args`, `kwargs` and `cancel` might be modified by :class:`~TaskModifier` instances.
"""
__slots__ = ('__function', '__args', '__kwargs', '__cancel')
def __init__(
self,
function: FUNCTION,
args: typing.Optional[typing.Iterable[typing.Any]] = None,
kwargs: typing.Optional[typing.Mapping[str, typing.Any]] = None,
cancel: typing.Optional[ACTION] = None,
) -> None:
"""Initialize new :class:`~TaskContext` instance."""
if args is None:
args = []
if kwargs is None:
kwargs = {}
self.__function: typing.Final[FUNCTION] = function
self.__args: typing.Final[typing.List[typing.Any]] = list(args)
self.__kwargs: typing.Final[typing.Dict[str, typing.Any]] = dict(kwargs)
self.__cancel: typing.Optional[ACTION] = cancel
@property
def args(self) -> typing.List[typing.Any]: # noqa: D401
"""Collection of positional arguments to call :attr:`~TaskContext.function` with."""
return self.__args
@property
def call(self) -> Call: # noqa: D401
"""New :class:`~Call` signature from the collected context."""
return Call(function=self.function, args=self.args, kwargs=self.kwargs)
@property
def cancel(self) -> typing.Optional[ACTION]: # noqa: D401
"""Function that can cancel function call in a graceful way."""
return self.__cancel
@cancel.setter
def cancel(self, value: typing.Optional[ACTION]) -> None:
"""Update `cancel` value."""
self.__cancel = value
@property
def function(self) -> FUNCTION: # noqa: D401
"""Function to call from this context."""
return self.__function
@property
def kwargs(self) -> typing.Dict[str, typing.Any]: # noqa: D401
"""Collection of keyword arguments to call :attr:`~TaskContext.function` with."""
return self.__kwargs
@property
def worker(self) -> TaskWorker: # noqa: D401
"""New :class:`~TaskWorker` from the collected context."""
return TaskWorker(call=self.call, cancel=self.cancel)
class TaskModifier(metaclass=abc.ABCMeta): # pylint: disable=too-few-public-methods
"""
Common interface for all :class:`~TaskContext` modifiers.
This is used to separate context modifiers from the actual function in :class:`~Task`.
"""
__slots__ = ()
@abc.abstractmethod
def apply(self, context: TaskContext) -> None:
"""Apply modification to the :class:`~TaskContext` instance."""
class InsertArgument(TaskModifier):
"""
Insert new positional argument.
Places `value` at `target` in :attr:`~TaskContext.args` or :attr:`~TaskContext.kwargs`.
:param value: Value of the argument to add.
:param target: Index of the positional or Key of the keyword argument to add.
"""
__slots__ = ('__value', '__target')
def __init__(self, value: typing.Any, target: typing.Union[int, str] = 0) -> None:
"""Initialize new :class:`~InsertArgument` instance."""
self.__value: typing.Final[typing.Any] = value
self.__target: typing.Final[typing.Union[int, str]] = target
@property
def target(self) -> typing.Union[int, str]: # noqa: D401
"""
Location for the new argument.
Might be :code:`int` index for positional arguments, or :code:`str` key for keyword arguments.
"""
return self.__target
@property
def value(self) -> typing.Any: # noqa: D401
"""Value of the new argument."""
return self.__value
def apply(self, context: TaskContext) -> None:
"""Apply modification to the :class:`~TaskContext` instance."""
if isinstance(self.target, int):
context.args.insert(self.target, self.value)
else:
context.kwargs[self.target] = self.value
class AddCancelCallback(TaskModifier):
"""
Add new callback to use for canceling a task.
:param callback: Callback function to use.
"""
__slots__ = ('__callback',)
def __init__(self, callback: ACTION) -> None:
"""Initialize new :class:`~AddCancelCallback` instance."""
self.__callback: typing.Final[ACTION] = callback
@property
def callback(self) -> ACTION: # noqa: D401
"""Callback function to use for canceling a task."""
return self.__callback
def apply(self, context: TaskContext) -> None:
"""Apply modification to the :class:`~TaskContext` instance."""
context.cancel = self.callback
# Ideas for extra modifiers:
# - AppendArgument(value) -> args.append(value)
# - EnsureArgument(target, value) -> kwargs.setdefault(target, value)
class Task:
"""
Decorator for functions and methods to convert them to tasks.
After decoration - functions and methods might still be called usual way. :meth:`~Task.worker` might be used to
create a worker for a task.
:param args: :class:`~TaskModifier` instances to modify a :class:`~TaskContext` before execution.
.. note::
Decorator can be called in both :code:`@Task` and :code:`@Task(...)` ways. Both will work.
"""
__slots__ = ('__function', '__modifiers')
@typing.overload
def __init__(self, *args: TaskModifier) -> None:
"""Initialize new :class:`~Task` instance."""
@typing.overload
def __init__(self, function: FUNCTION, /, *args: TaskModifier) -> None:
"""Initialize new :class:`~Task` instance."""
def __init__(self, *args: typing.Union[FUNCTION, TaskModifier]) -> None:
"""Initialize new :class:`~Task` instance."""
function: typing.Optional[FUNCTION] = None
modifiers: typing.List[TaskModifier] = []
arg: typing.Union[FUNCTION, TaskModifier]
for arg in args:
if isinstance(arg, TaskModifier):
modifiers.append(arg)
elif function is None:
function = arg
else:
raise TypeError('Task arguments must be TaskModifiers and at most one function')
self.__function: typing.Optional[FUNCTION] = function
self.__modifiers: typing.Final[typing.Tuple[TaskModifier, ...]] = tuple(modifiers)
@property
def function(self) -> FUNCTION: # noqa: D401
"""Function that is used for this task."""
if self.__function is None:
raise AttributeError('function was not wrapped yet')
return self.__function
@property
def modifiers(self) -> typing.Tuple[TaskModifier, ...]: # noqa: D401
"""Collection of modifiers for a :class:`~TaskContext`."""
return self.__modifiers
def __context(
self,
args: typing.Optional[typing.Iterable[typing.Any]] = None,
kwargs: typing.Optional[typing.Mapping[str, typing.Any]] = None,
) -> TaskContext:
"""Prepare a context for execution or converting into a task."""
modifier: TaskModifier
result: typing.Final[TaskContext] = TaskContext(function=self.function, args=args, kwargs=kwargs)
for modifier in self.modifiers:
modifier.apply(context=result)
return result
def __call__(self, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
"""
Execute a function in a usual way.
.. note::
This method might be also called when the decorator is called in a :code:`@Task(...)` way to finalize
function/method wrapping.
"""
if self.__function is None:
if (len(args) != 1) or kwargs:
raise TypeError('function must be the only provided argument for wrapping.')
self.__function = args[0]
return self
return self.__context(args=args, kwargs=kwargs).call()
def __get__(self, instance: typing.Any, owner: typing.Any) -> 'Task':
"""Prepare a bound task for a methods called for a specific instances."""
if instance is None:
return self
return Task(self.function, *self.modifiers, InsertArgument(instance))
def worker(self, *args: typing.Any, **kwargs: typing.Any) -> TaskWorker:
"""Prepare worker to execute this task."""
return self.__context(args=args, kwargs=kwargs).worker
class TaskManager(QtCore.QObject):
"""Central manager for executing tasks."""
sig_empty = QtCore.Signal()
sig_start = QtCore.Signal(TaskWorker)
sig_succeeded = QtCore.Signal(TaskWorker)
sig_canceled = QtCore.Signal(TaskWorker)
sig_failed = QtCore.Signal(TaskWorker)
sig_done = QtCore.Signal(TaskWorker)
def __init__(self) -> None:
"""Initialize new :class:`~TaskManager` instance."""
super().__init__()
self.__cancel: bool = False
self.__workers: typing.Final[typing.Set[TaskWorker]] = set()
def __start(self, worker: TaskWorker) -> None:
"""Process start of a worker."""
self.__workers.add(worker)
if self.__cancel:
with contextlib.suppress(TypeError):
worker.cancel()
self.sig_start.emit(worker)
def __succeeded(self, worker: TaskWorker) -> None:
"""Process success of a worker."""
if worker not in self.__workers:
return
self.sig_succeeded.emit(worker)
def __canceled(self, worker: TaskWorker) -> None:
"""Process cancel of a worker."""
if worker not in self.__workers:
return
self.sig_canceled.emit(worker)
def __failed(self, worker: TaskWorker) -> None:
"""Process fail of a worker."""
if worker not in self.__workers:
return
self.sig_failed.emit(worker)
def __done(self, worker: TaskWorker) -> None:
"""Process worker finished execution."""
try:
self.__workers.remove(worker)
except LookupError:
return
self.sig_done.emit(worker)
if not self.__workers:
self.sig_empty.emit()
def _register(self, worker: TaskWorker) -> None:
"""Register worker when it starts its execution."""
if worker in self.__workers:
return
self.__start(worker)
worker.signals.sig_succeeded.connect(lambda _: self.__succeeded(worker))
worker.signals.sig_canceled.connect(lambda _: self.__canceled(worker))
worker.signals.sig_failed.connect(lambda _: self.__failed(worker))
worker.signals.sig_done.connect(lambda _: self.__done(worker))
def cancel_all(self) -> None:
"""Cancel all running and future workers."""
self.__cancel = True
workers: typing.Sequence[TaskWorker] = list(self.__workers)
if not workers:
self.sig_empty.emit()
return
worker: TaskWorker
for worker in workers:
with contextlib.suppress(TypeError):
worker.cancel()
def reset(self) -> None:
"""
Reset task processing.
This might be useful to allow executing new tasks after calling :meth:`~TaskManager.cancel_all`.
"""
self.__cancel = False
class CancelContext:
"""
Simplistic context with a canceled flag.
:param canceled: Initial canceled state.
"""
__slots__ = ('__canceled',)
def __init__(self, canceled: bool = False) -> None:
"""Initialize new :class:`~CancelContext` instance."""
self.__canceled: bool = canceled
@property
def canceled(self) -> bool: # noqa: D401
"""Task was canceled."""
return self.__canceled
def abort_if_canceled(self) -> None:
"""Break execution of a task if it was canceled."""
if self.__canceled:
raise TaskCanceledError()
def cancel(self) -> None:
"""Cancel task execution."""
self.__canceled = True
class AddCancelContext(TaskModifier):
"""
Add new :class:`~CancelContext` to the :class:`~TaskContext`.
This adds new :class:`~CancelContext` instance to the function kwargs and modifies :attr:`~TaskContext.cancel` to
use the context.
:param argument: Name of the keyword argument to add context to.
.. warning::
Context will be provided if only it was not provided before (with a regular keyword argument).
Otherwise - it will reuse provided value.
"""
__slots__ = ('__argument',)
def __init__(self, argument: str = 'context') -> None:
"""Initialize new :class:`~AddCancelContext` instance."""
self.__argument: typing.Final[str] = argument
@property
def argument(self) -> str: # noqa: D401
"""Name of the argument to which context will be added."""
return self.__argument
def apply(self, context: TaskContext) -> None:
"""Apply modification to the :class:`~TaskContext` instance."""
if self.argument not in context.kwargs:
context.kwargs[self.argument] = CancelContext()
context.cancel = context.kwargs[self.argument].cancel
class ThreadPoolSingleton(singletons.Singleton[QtCore.QThreadPool]): # pylint: disable=too-few-public-methods
"""Singleton instance of a :class:`~qtpy.QtCore.QThreadPool`."""
__slots__ = ()
def _prepare(self) -> QtCore.QThreadPool:
"""Initialize singleton instance."""
result: QtCore.QThreadPool = QtCore.QThreadPool()
result.setMaxThreadCount(max(8, result.maxThreadCount()))
return result
def _release(self) -> None:
"""Destroy singleton instance when :meth:`~Singleton.reset` is called."""
self.instance.clear()
self.instance.waitForDone()
MANAGER: typing.Final[singletons.Singleton[TaskManager]] = singletons.SingleInstanceOf(TaskManager)
THREAD_POOL: typing.Final[singletons.Singleton[QtCore.QThreadPool]] = ThreadPoolSingleton()
def teardown() -> None:
"""Finish all background workers."""
MANAGER.instance.cancel_all()
THREAD_POOL.reset()