Source code for schrodinger.ui.qt.jobwidgets

"""
This module provides job widgets compatible with all Schrodinger python
application frameworks.

Copyright Schrodinger, LLC. All rights reserved.
"""
from typing import Optional

from schrodinger import get_maestro
from schrodinger.application.job_monitor import job_monitor_gui
from schrodinger.infra import jobhub
from schrodinger.models import mappers
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.tasks import taskmanager
from schrodinger.ui.qt.standard.icons import icons
from schrodinger.ui.qt.standard_widgets import buttons
from schrodinger.utils import mmutil

maestro = get_maestro()

SPINNER_ICON_BASE = ":/schrodinger/ui/qt/icons_dir/spinner/"
MAX_SPINNER_ICON_NUM = 9

JobIconState = jobhub.ProjectLatestJobIconTracker.JobIconState


class JobStatusButton(mappers.TargetMixin, buttons.FlatButton):
    """
    A button with a spinner icon that spins upon the parent widget launching
    a job. It continues to spin for up to 30 seconds after the job has started,
    or until the job finishes or fails. Failures result in an error icon.

    Clicking this tool button will launch the job monitor panel.

    This button also implements `targetSetValue` to support mapping to a task or
    taskmanager.

    Note: This button will not show/hide properly if added to a QToolBar since
    this is not a QAction. It works well inside a layout.
    """

    def __init__(self,
                 parent: Optional[QtWidgets.QWidget] = None,
                 viewname: Optional[str] = None):
        """
        :ivar viewname: View name to use for job launching and monitoring.
            If not specified, use setViewname() to set it at a later time.
        """
        super().__init__(parent=parent)
        self.setToolTip('Click to display the job monitor')
        self._updateVisibility()
        self._setUpSpinnerIcons()
        self._icon_state_tracker = None
        self._viewname = None
        if viewname:
            self.setViewname(viewname)

        self.pressed.connect(self._showJobMonitor)

    def _setUpSpinnerIcons(self):
        self._icons = [
            QtGui.QIcon(f'{SPINNER_ICON_BASE}{num}.png')
            for num in range(1, MAX_SPINNER_ICON_NUM)
        ]
        self._standard_icon = QtGui.QIcon(f'{SPINNER_ICON_BASE}0.png')
        self._error_icon = QtGui.QIcon(icons.EXCLAMATION_ERROR_LB)
        self._cur_icon_idx = 0
        self._num_icons = len(self._icons)
        self.setIcon(self._standard_icon)
        self._spin_timer = QtCore.QTimer(self)
        self._spin_timer.timeout.connect(self._advanceSpinner)

    def _setUpIconStateTracker(self):
        """
        Set up the icon state tracker. Called whenever the viewname is set.
        """
        if self._icon_state_tracker is not None:
            self._icon_state_tracker.jobIconStateChanged.disconnect(
                self._updateSpinnerIconState)
        if self._viewname is None:
            self._icon_state_tracker = None
            return
        self._icon_state_tracker = jobhub.ProjectLatestJobIconTracker(
            self, jobhub.get_job_manager(), self._viewname)
        self._icon_state_tracker.jobIconStateChanged.connect(
            self._updateSpinnerIconState)
        self._icon_state_tracker.jobIconStateChanged.connect(
            self._updateVisibility)

    # ==========================================================================
    # Public API
    # ==========================================================================

    def setViewname(self, viewname: Optional[str] = None):
        """
        Set the viewname and instantiate the job icon state tracker.
        """
        old_viewname = self._viewname
        self._viewname = viewname
        if old_viewname != viewname and viewname is not None:
            self._setUpIconStateTracker()

    # ==========================================================================
    # TargetMixin methods
    # ==========================================================================

    def targetSetValue(self, task_or_taskman):
        if isinstance(task_or_taskman, taskmanager.TaskManager):
            task = task_or_taskman.nextTask()
        else:
            task = task_or_taskman
        self.setViewname(task.job_config.viewname)

    # ==========================================================================
    # Slots
    # ==========================================================================

    def _updateSpinnerIconState(self, icon_state: JobIconState):
        """
        Called when the icon state tracker emits jobIconStateChanged.

        :param icon_state: The current icon state.
        """
        if icon_state == JobIconState.Active:
            self._startSpinner()
            return
        self._stopSpinner()
        if icon_state == JobIconState.Fail:
            self.setIcon(self._error_icon)

    def _updateVisibility(self, icon_state: Optional[JobIconState] = None):
        """
        Update this button's visibility when the job icon state changes.
        """
        visible = icon_state in (JobIconState.Active, JobIconState.Fail)
        self.setVisible(visible)

    def _showJobMonitor(self):
        """
        Show the correct job monitor and display the current job if a job is
        currently being tracked. Will not display anything outside of maestro
        if JOB_SERVER is off.

        Called when this button is pressed.
        """
        if self._icon_state_tracker is None:
            return
        job_id = self._icon_state_tracker.getJobId()
        if mmutil.feature_flag_is_enabled(mmutil.JOB_SERVER):
            # Show new job monitor
            job_monitor = job_monitor_gui.panel()
            if job_id:
                job_monitor.selectJob(job_id)
        elif maestro:
            # Show legacy job monitor only in maestro
            maestro.command('showpanel monitor')
            if job_id:
                maestro.command(f'energymonitor {job_id}')

    # ==========================================================================
    # Private implementation methods
    # ==========================================================================

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

    def _stopSpinner(self):
        if self._spin_timer.isActive():
            self._spin_timer.stop()
        self.setIcon(self._standard_icon)

    def _advanceSpinner(self):
        self._cur_icon_idx = (self._cur_icon_idx + 1) % self._num_icons
        self.setIcon(self._icons[self._cur_icon_idx])