Source code for schrodinger.ui.qt.tasks.taskwidgets

import os

from schrodinger.models import advanced_mappers
from schrodinger.models import mappers
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.Qt.QtCore import Qt
from schrodinger.tasks import gui
from schrodinger.tasks import jobtasks
from schrodinger.tasks import taskmanager
from schrodinger.tasks import tasks
from schrodinger.ui.qt import basewidgets
from schrodinger.ui.qt import jobwidgets
from schrodinger.ui.qt import utils
from schrodinger.ui.qt.tasks import taskbarwidgets
from schrodinger.ui.qt.tasks.taskbarwidgets import TaskWidgetMixin

DEFAULT_START_LATENCY = 750
IMAGE_PATH = os.path.join(":/schrodinger/ui/qt/icons_dir/spinner/")
ICON_PATH_TEMPLATE = os.path.join(IMAGE_PATH, "{}.png")
SPIN_ICONS = None
ERROR_ICON = None


def get_spin_icons():
    global SPIN_ICONS
    if SPIN_ICONS is None:
        SPIN_ICONS = [
            QtGui.QIcon(ICON_PATH_TEMPLATE.format(num)) for num in range(0, 9)
        ]
    return SPIN_ICONS


def get_error_icon():
    global ERROR_ICON
    ERROR_ICON = QtGui.QIcon(ICON_PATH_TEMPLATE.format('error'))
    return ERROR_ICON


class _SpinnerMixin:
    running_pics = NotImplemented
    stop_pic = NotImplemented
    error_pic = NotImplemented

    def setImage(self, image):
        raise NotImplementedError

    def __init__(self, *args, **kwargs):
        self.current_num = 0
        self.timer = QtCore.QTimer()
        self.timer.timeout.connect(self._advanceSpinner)
        super().__init__(*args, **kwargs)

    def _advanceSpinner(self):
        pic = self.running_pics[self.current_num]
        self.setImage(pic)
        self.current_num += 1
        if self.current_num == len(self.running_pics):
            self.current_num = 0

    def _startSpinner(self):
        if not self.timer.isActive():
            self.timer.start(250)

    def _stopSpinner(self):
        self.timer.stop()
        task = self.model
        if task is None:
            return
        if task.status is task.FAILED:
            self.setImage(self.error_pic)
        else:
            self.setImage(self.stop_pic)

    def defineMappings(self):
        mappings = super().defineMappings()
        M = self.model_class
        target = mappers.TargetSpec(slot=self.updateSpinnerState)
        mappings += [(target, M.status)]
        return mappings

    def updateSpinnerState(self):
        status = self.model.status
        if status != tasks.Status.RUNNING:
            self._stopSpinner()
        else:
            self._startSpinner()


class SpinnerLabel(_SpinnerMixin, TaskWidgetMixin, QtWidgets.QLabel):
    clicked = QtCore.pyqtSignal()

    def __init__(self, *args, **kwargs):
        height = 16
        icons = get_spin_icons()
        self.running_pics = [icon.pixmap(height, height) for icon in icons[1:]]
        self.stop_pic = icons[0].pixmap(height, height)
        self.error_pic = get_error_icon().pixmap(height, height)
        super().__init__(*args, **kwargs)

    def mousePressEvent(self, event):
        QtWidgets.QLabel.mousePressEvent(self, event)
        self.clicked.emit()

    def setImage(self, image):
        self.setPixmap(image)


class ProgressBar(TaskWidgetMixin, QtWidgets.QProgressBar):
    """
    Progress bar for a running task. For AbstractTask, this monitors the
    task.status, task.progress, and task.max_progress and updates accordingly.
    For TaskManager, this monitors taskman.status only.

    Set max_progress to 0 for an indeterminate progress bar. Both progress and
    max_progress may be adjusted while the task is running.

    Determinate progress is only supported for tasks. For a taskmanager, the
    progress dialog will run as long as any task in the taskmanager is running.

    """

    def __init__(self, *args, **kwargs):
        self._is_running = False
        super().__init__(*args, **kwargs)
        self.setMinimum(0)
        self.setMaximum(1)
        self._stop()

    def defineMappings(self):
        M = self.model_class
        status_target = mappers.TargetSpec(slot=self.updateStatus)
        mappings = [(status_target, M.status)]
        if M is tasks.AbstractTask:
            progress_target = mappers.TargetSpec(slot=self.updateProgress)
            mappings.append((progress_target, M.progress))
            mappings.append((progress_target, M.max_progress))
        return mappings

    def setModel(self, model):
        if model is None:
            self._stop()
        super().setModel(model)

    def updateStatus(self):
        M = self.model_class
        status = self.model.status
        if status == tasks.Status.RUNNING:
            if M is tasks.AbstractTask:
                task = self.model
                self.setMaximum(task.max_progress)
            if not self.isRunning():
                self._start()
        else:
            if self.isRunning():
                self._stop()

    def updateProgress(self):
        """
        Updates the numerical progress. Not called for TaskManager.
        """
        task = self.model
        self.setMaximum(task.max_progress)
        self.setValue(task.progress)
        self._setDeterminate(task.max_progress > 0)

    def _setDeterminate(self, is_determinate):
        self.setTextVisible(is_determinate)

    def isRunning(self):
        return self._is_running

    def _start(self):
        M = self.model_class
        if M is taskmanager.TaskManager:
            self.setMaximum(0)
        self._is_running = True

    def _stop(self):
        M = self.model_class
        if (M is taskmanager.TaskManager or self.model is None or
                self.model.max_progress == 0):
            # This is how you stop an indeterminate progress bar
            self.setMaximum(1)
        self._is_running = False


class ProgressDialog(advanced_mappers.MultiModelClassMapperMixin,
                     basewidgets.BaseWidget):
    """
    A basic modal progress dialog that shows itself whenever its model task is
    running and hides itself whenever it's not running. The dialog contains a
    ProgressBar instance and a text label, which can be set with setLabelText.

    Calling setTask with a running task will result in the dialog immediately
    being shown.
    """
    model_classes = (tasks.AbstractTask, taskmanager.TaskManager)
    SHOW_AS_WINDOW = True

    def initSetUp(self):
        super().initSetUp()
        self._is_running = False
        self.progress_bar = ProgressBar()
        self.main_lbl = QtWidgets.QLabel()

    def initLayOut(self):
        super().initLayOut()
        self.main_layout.addWidget(self.main_lbl)
        self.main_lbl.setAlignment(Qt.AlignHCenter)
        self.main_layout.addWidget(self.progress_bar)
        self.main_layout.setContentsMargins(9, 9, 9, 9)
        self.main_layout.setSpacing(9)

    def setTask(self, task):
        self.setModel(task)

    def setModel(self, model):
        if model is None:
            self._stop()
        super().setModel(model)

    def setLabelText(self, text):
        self.main_lbl.setText(text)

    def defineMappings(self):
        M = self.model_class
        status_target = mappers.TargetSpec(slot=self.updateStatus)
        return [(self.progress_bar, M), (status_target, M.status)]

    def updateStatus(self):
        task = self.model
        if task.isRunning():
            self._start()
        else:
            self._stop()

    def isRunning(self):
        return self._is_running

    def _start(self):
        self._is_running = True
        self.run(modal=True)

    def _stop(self):
        self._is_running = False
        self.hide()


class KillableProgressDialog(ProgressDialog):
    """
    A `ProgressDialog` with a "Cancel" button that kills the running task.
    """

    def initSetUp(self):
        super().initSetUp()
        self.kill_task_btn = QtWidgets.QPushButton('Cancel', self)
        self.kill_task_btn.clicked.connect(self._onKillTaskClicked)

    def initLayOut(self):
        super().initLayOut()
        self.bottom_middle_layout.addWidget(self.kill_task_btn)

    def _onKillTaskClicked(self):
        task = self.model
        task.kill()


class StatusLabel(TaskWidgetMixin, QtWidgets.QLabel):

    def defineMappings(self):
        M = self.model_class
        target = mappers.TargetSpec(setter=self.setStatus)
        return [(target, M.status)]

    def setStatus(self, value):
        self.setText(value.name)


class StartButton(utils.ButtonAcceptsFocusMixin, TaskWidgetMixin,
                  QtWidgets.QPushButton):
    """
    A button responsible for launching a task.

    :ivar aboutToStartTask: a signal emitted just before the task is launched;
            contains the task instance as an argument
    :vartype aboutToStartTask: QtCore.pyqtSignal
    """

    aboutToStartTask = QtCore.pyqtSignal(tasks.AbstractTask)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.clicked.connect(self.onClicked)
        if not self.text():
            self.setText('Start')

    def updateEnabled(self, value=None):
        startable = self.model is not None and self.model.isStartable()
        self.setEnabled(startable)

    def defineMappings(self):
        M = self.model_class
        target = mappers.TargetSpec(slot=self.updateEnabled)
        return [(target, M.status)]

    def onClicked(self):
        M = self.model_class
        if M is tasks.AbstractTask:
            task = self.model
        elif M is taskmanager.TaskManager:
            task = self.model.nextTask()
        self.aboutToStartTask.emit(task)
        gui.start_task(task, self.window())


class SpinnerStartButton(_SpinnerMixin, StartButton):

    def __init__(self, *args, **kwargs):
        self.running_pics = get_spin_icons()
        self.stop_pic = self.running_pics[0]
        self.error_pic = get_error_icon()
        super().__init__(*args, **kwargs)

    def setImage(self, image):
        self.setIcon(image)


class SettingsButton(QtWidgets.QToolButton):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.setObjectName("af2SettingsButton")
        self.setPopupMode(QtWidgets.QToolButton.MenuButtonPopup)
        self.setIcon(QtGui.QIcon(':/icons/small_settings.png'))

    def setActions(self, text_to_slots):
        """
        Set the actions to show for the settings menu that popups when pressing
        the down arrow on the button.

        :param text_to_slots: A list of tuples mapping the desired menu item text
            with the function that should be called when the item is selected.
            If tuple is `None`, then a separator will be added instead and
            the text will be ignored.
        :type  text_to_slots: list[tuple(str, callable) or None]
        """
        menu = QtWidgets.QMenu()
        for text_and_slot in text_to_slots:
            if text_and_slot is None:
                menu.addSeparator()
            else:
                menu.addAction(*text_and_slot)
        self.setMenu(menu)


class AbstractTaskBar(mappers.MapperMixin, basewidgets.BaseWidget):
    """
    Base class for all taskbars. To create a fully customized taskbar,
    subclass this class and make sure that startRequested is emitted
    whenever a task should be started.
    """
    model_class = taskmanager.TaskManager
    startRequested = QtCore.pyqtSignal()

    def setConfigDialog(self, config_dialog):
        """
        Set a config dialog for this taskbar.

        Subclasses should define this if they support a config dialog.
        """
        raise NotImplementedError

    def makeInitialModel(self):
        """
        @overrides: mappers.MapperMixin
        """
        return None

    def _getCurrentTask(self):
        if isinstance(self.model, taskmanager.TaskManager):
            return self.model.nextTask()
        else:
            return self.model

    def setSettingsMenuActions(self, text_to_slots):
        """
        Set the actions to show for the settings menu shown from the settings
        tool button.

        Subclasses should define this if they have a settings button.

        :param text_to_slots: A list of tuples mapping the desired menu item
            text with the function that should be called when the item is
            selected.  If the slot is `None`, then a separator will be added
            instead and the text will be ignored.
        :type  text_to_slots: list[tuple(str, callable) or None]
        """
        raise NotImplementedError


class TaskBar(AbstractTaskBar):

    def initSetUp(self):
        super().initSetUp()
        self.name_le = taskbarwidgets.NameLineEdit()
        self.task_name_lbl = QtWidgets.QLabel("Job name: ")
        self.start_btn = utils.AcceptsFocusPushButton("Run")
        self.spinner_label = SpinnerLabel()
        self.settings_btn = SettingsButton()
        self.config_dialog = None
        # We connect signals later so subclasses can override these
        # widgets with their own.

    def initFinalize(self):
        super().initFinalize()
        self._connectSignals()
        if not isinstance(self.start_btn, utils.ButtonAcceptsFocusMixin):
            msg = (f'Expected `self.start_btn` class'
                   f'{self.start_btn.__class__.__name__} to inherit '
                   'ButtonAcceptsFocusMixin')
            raise TypeError(msg)

    def _connectSignals(self):
        self.start_btn.clicked.connect(self.startRequested)

    def setConfigDialog(self, config_dialog):
        config_dialog.startRequested.connect(self.startRequested)
        self.settings_btn.clicked.connect(self.showConfigDialog)
        self.config_dialog = config_dialog

    def showConfigDialog(self):
        if self.config_dialog is None:
            raise RuntimeError("No config dialog has been set for this taskbar")
        self.config_dialog.setParent(self)
        self.config_dialog.setModel(self._getCurrentTask().job_config)
        self.config_dialog.run(modal=True)

    def initLayOut(self):
        super().initLayOut()
        hlayout = QtWidgets.QHBoxLayout()
        hlayout.addWidget(self.task_name_lbl)
        hlayout.addWidget(self.name_le)
        hlayout.addWidget(self.settings_btn)
        hlayout.addWidget(self.spinner_label)
        hlayout.addWidget(self.start_btn)

        self.main_layout.addLayout(hlayout)

    def setModel(self, model):
        super().setModel(model)
        self._updateSettingsBtnVisibility()

    def setSettingsMenuActions(self, text_to_slots):
        """
        Set the actions to show for the settings menu shown from the settings
        tool button.

        :param text_to_slots: A list of tuples mapping the desired menu item
            text with the function that should be called when the item is
            selected.  If the slot is `None`, then a separator will be added
            instead and the text will be ignored.
        :type  text_to_slots: list[tuple(str, callable) or None]
        """
        self.settings_btn.setActions(text_to_slots)

    def _updateSettingsBtnVisibility(self):
        """
        Show the settings button if the task is a jobtask, otherwise hide it.
        """
        task = self._getCurrentTask()
        should_show_settings = jobtasks.is_jobtask(task)
        self.settings_btn.setVisible(should_show_settings)

    def defineMappings(self):
        M = self.model_class
        mappings = [(self.spinner_label, M),
                    (self.name_le, M),
                    (self.settings_btn, M),
                    (self.start_btn, M)] # yapf: disable
        if M is taskmanager.TaskManager:
            t = mappers.TargetSpec(slot=self._updateSettingsBtnVisibility)
            mappings.append((t, M.TaskClass))
        return mappings


class JobTaskBar(TaskBar):

    def initSetUp(self):
        super().initSetUp()
        self.spinner_label = jobwidgets.JobStatusButton(parent=self)