ayon-core/openpype/tools/context_dialog/window.py

import os
import json

from qtpy import QtWidgets, QtCore, QtGui

from openpype import style
from openpype.pipeline import AvalonMongoDB
from openpype.tools.utils.lib import center_window, get_openpype_qt_app
from openpype.tools.utils.assets_widget import SingleSelectAssetsWidget
from openpype.tools.utils.constants import (
    PROJECT_NAME_ROLE
)
from openpype.tools.utils.tasks_widget import TasksWidget
from openpype.tools.utils.models import (
    ProjectModel,
    ProjectSortFilterProxy
)


class ContextDialog(QtWidgets.QDialog):
    """Dialog to select a context.

    Context has 3 parts:
    - Project
    - Aseet
    - Task

    It is possible to predefine project and asset. In that case their widgets
    will have passed preselected values and will be disabled.
    """
    def __init__(self, parent=None):
        super(ContextDialog, self).__init__(parent)

        self.setWindowTitle("Select Context")
        self.setWindowIcon(QtGui.QIcon(style.app_icon_path()))

        # Enable minimize and maximize for app
        window_flags = QtCore.Qt.Window
        if not parent:
            window_flags |= QtCore.Qt.WindowStaysOnTopHint
        self.setWindowFlags(window_flags)
        self.setFocusPolicy(QtCore.Qt.StrongFocus)

        dbcon = AvalonMongoDB()

        # UI initialization
        main_splitter = QtWidgets.QSplitter(self)

        # Left side widget contains project combobox and asset widget
        left_side_widget = QtWidgets.QWidget(main_splitter)

        project_combobox = QtWidgets.QComboBox(left_side_widget)
        # Styled delegate to propagate stylessheet
        project_delegate = QtWidgets.QStyledItemDelegate(project_combobox)
        project_combobox.setItemDelegate(project_delegate)
        # Project model with only active projects without default item
        project_model = ProjectModel(
            dbcon,
            only_active=True,
            add_default_project=False
        )
        # Sorting proxy model
        project_proxy = ProjectSortFilterProxy()
        project_proxy.setSourceModel(project_model)
        project_combobox.setModel(project_proxy)

        # Assets widget
        assets_widget = SingleSelectAssetsWidget(
            dbcon, parent=left_side_widget
        )

        left_side_layout = QtWidgets.QVBoxLayout(left_side_widget)
        left_side_layout.setContentsMargins(0, 0, 0, 0)
        left_side_layout.addWidget(project_combobox)
        left_side_layout.addWidget(assets_widget)

        # Right side of window contains only tasks
        tasks_widget = TasksWidget(dbcon, main_splitter)

        # Add widgets to main splitter
        main_splitter.addWidget(left_side_widget)
        main_splitter.addWidget(tasks_widget)

        # Set stretch of both sides
        main_splitter.setStretchFactor(0, 7)
        main_splitter.setStretchFactor(1, 3)

        # Add confimation button to bottom right
        ok_btn = QtWidgets.QPushButton("OK", self)

        buttons_layout = QtWidgets.QHBoxLayout()
        buttons_layout.setContentsMargins(0, 0, 0, 0)
        buttons_layout.addStretch(1)
        buttons_layout.addWidget(ok_btn, 0)

        main_layout = QtWidgets.QVBoxLayout(self)
        main_layout.addWidget(main_splitter, 1)
        main_layout.addLayout(buttons_layout, 0)

        # Timer which will trigger asset refresh
        # - this is needed because asset widget triggers
        #   finished refresh before hides spin box so we need to trigger
        #   refreshing in small offset if we want re-refresh asset widget
        assets_timer = QtCore.QTimer()
        assets_timer.setInterval(50)
        assets_timer.setSingleShot(True)

        assets_timer.timeout.connect(self._on_asset_refresh_timer)

        project_combobox.currentIndexChanged.connect(
            self._on_project_combo_change
        )
        assets_widget.selection_changed.connect(self._on_asset_change)
        assets_widget.refresh_triggered.connect(self._on_asset_refresh_trigger)
        assets_widget.refreshed.connect(self._on_asset_widget_refresh_finished)
        tasks_widget.task_changed.connect(self._on_task_change)
        ok_btn.clicked.connect(self._on_ok_click)

        self._dbcon = dbcon

        self._project_combobox = project_combobox
        self._project_model = project_model
        self._project_proxy = project_proxy
        self._project_delegate = project_delegate

        self._assets_widget = assets_widget

        self._tasks_widget = tasks_widget

        self._ok_btn = ok_btn

        self._strict = False

        # Values set by `set_context` method
        self._set_context_project = None
        self._set_context_asset = None

        # Requirements for asset widget refresh
        self._assets_timer = assets_timer
        self._rerefresh_assets = True
        self._assets_refreshing = False

        # Set stylehseet and resize window on first show
        self._first_show = True

        # Helper attributes for handling of refresh
        self._ignore_value_changes = False
        self._refresh_on_next_show = True

        # Output of dialog
        self._context_to_store = {
            "project": None,
            "asset": None,
            "task": None
        }

    def closeEvent(self, event):
        """Ignore close event if is in strict state and context is not done."""
        if self._strict and not self._ok_btn.isEnabled():
            event.ignore()
            return

        if self._strict:
            self._confirm_values()
        super(ContextDialog, self).closeEvent(event)

    def set_strict(self, strict):
        """Change strictness of dialog."""
        self._strict = strict
        self._validate_strict()

    def _set_refresh_on_next_show(self):
        """Refresh will be called on next showEvent.

        If window is already visible then just execute refresh.
        """
        self._refresh_on_next_show = True
        if self.isVisible():
            self.refresh()

    def _refresh_assets(self):
        """Trigger refreshing of asset widget.

        This will set mart to rerefresh asset when current refreshing is done
        or do it immidietely if asset widget is not refreshing at the time.
        """
        if self._assets_refreshing:
            self._rerefresh_assets = True
        else:
            self._on_asset_refresh_timer()

    def showEvent(self, event):
        """Override show event to do some callbacks."""
        super(ContextDialog, self).showEvent(event)
        if self._first_show:
            self._first_show = False
            # Set stylesheet and resize
            self.setStyleSheet(style.load_stylesheet())
            self.resize(600, 700)
            center_window(self)

        if self._refresh_on_next_show:
            self.refresh()

    def refresh(self):
        """Refresh all widget one by one.

        When asset refresh is triggered we have to wait when is done so
        this method continues with `_on_asset_widget_refresh_finished`.
        """
        # Change state of refreshing (no matter how refresh was called)
        self._refresh_on_next_show = False

        # Ignore changes of combobox and asset widget
        self._ignore_value_changes = True

        # Get current project name to be able set it afterwards
        select_project_name = self._dbcon.Session.get("AVALON_PROJECT")
        # Trigger project refresh
        self._project_model.refresh()
        # Sort projects
        self._project_proxy.sort(0)

        # Disable combobox if project was passed to `set_context`
        if self._set_context_project:
            select_project_name = self._set_context_project
            self._project_combobox.setEnabled(False)
        else:
            # Find new project to select
            self._project_combobox.setEnabled(True)
            if (
                select_project_name is None
                and self._project_proxy.rowCount() > 0
            ):
                index = self._project_proxy.index(0, 0)
                select_project_name = index.data(PROJECT_NAME_ROLE)

        self._ignore_value_changes = False

        idx = self._project_combobox.findText(select_project_name)
        if idx >= 0:
            self._project_combobox.setCurrentIndex(idx)
        self._dbcon.Session["AVALON_PROJECT"] = (
            self._project_combobox.currentText()
        )

        # Trigger asset refresh
        self._refresh_assets()

    def _on_asset_refresh_timer(self):
        """This is only way how to trigger refresh asset widget.

        Use `_refresh_assets` method to refresh asset widget.
        """
        self._assets_widget.refresh()

    def _on_asset_widget_refresh_finished(self):
        """Catch when asset widget finished refreshing."""
        # If should refresh again then skip all other callbacks and trigger
        #   assets timer directly.
        self._assets_refreshing = False
        if self._rerefresh_assets:
            self._rerefresh_assets = False
            self._assets_timer.start()
            return

        self._ignore_value_changes = True
        if self._set_context_asset:
            self._dbcon.Session["AVALON_ASSET"] = self._set_context_asset
            self._assets_widget.setEnabled(False)
            self._assets_widget.select_assets(self._set_context_asset)
            self._set_asset_to_tasks_widget()
        else:
            self._assets_widget.setEnabled(True)
            self._assets_widget.set_current_asset_btn_visibility(False)

        # Refresh tasks
        self._tasks_widget.refresh()

        self._ignore_value_changes = False

        self._validate_strict()

    def _on_project_combo_change(self):
        if self._ignore_value_changes:
            return
        project_name = self._project_combobox.currentText()

        if self._dbcon.Session.get("AVALON_PROJECT") == project_name:
            return

        self._dbcon.Session["AVALON_PROJECT"] = project_name

        self._refresh_assets()
        self._validate_strict()

    def _on_asset_refresh_trigger(self):
        self._assets_refreshing = True
        self._on_asset_change()

    def _on_asset_change(self):
        """Selected assets have changed"""
        if self._ignore_value_changes:
            return
        self._set_asset_to_tasks_widget()

    def _on_task_change(self):
        self._validate_strict()

    def _set_asset_to_tasks_widget(self):
        asset_id = self._assets_widget.get_selected_asset_id()

        self._tasks_widget.set_asset_id(asset_id)

    def _confirm_values(self):
        """Store values to output."""
        self._context_to_store["project"] = self.get_selected_project()
        self._context_to_store["asset"] = self.get_selected_asset()
        self._context_to_store["task"] = self.get_selected_task()

    def _on_ok_click(self):
        # Store values to output
        self._confirm_values()
        # Close dialog
        self.accept()

    def get_selected_project(self):
        """Get selected project."""
        return self._project_combobox.currentText()

    def get_selected_asset(self):
        """Currently selected asset in asset widget."""
        return self._assets_widget.get_selected_asset_name()

    def get_selected_task(self):
        """Currently selected task."""
        return self._tasks_widget.get_selected_task_name()

    def _validate_strict(self):
        if not self._strict:
            if not self._ok_btn.isEnabled():
                self._ok_btn.setEnabled(True)
            return

        enabled = True
        if not self._set_context_project and not self.get_selected_project():
            enabled = False
        elif not self._set_context_asset and not self.get_selected_asset():
            enabled = False
        elif not self.get_selected_task():
            enabled = False
        self._ok_btn.setEnabled(enabled)

    def set_context(self, project_name=None, asset_name=None):
        """Set context which will be used and locked in dialog."""
        if project_name is None:
            asset_name = None

        self._set_context_project = project_name
        self._set_context_asset = asset_name

        self._context_to_store["project"] = project_name
        self._context_to_store["asset"] = asset_name

        self._set_refresh_on_next_show()

    def get_context(self):
        """Result of dialog."""
        return self._context_to_store


def main(
    path_to_store,
    project_name=None,
    asset_name=None,
    strict=True
):
    # Run Qt application
    app = get_openpype_qt_app()
    window = ContextDialog()
    window.set_strict(strict)
    window.set_context(project_name, asset_name)
    window.show()
    app.exec_()

    # Get result from window
    data = window.get_context()

    # Make sure json filepath directory exists
    file_dir = os.path.dirname(path_to_store)
    if not os.path.exists(file_dir):
        os.makedirs(file_dir)

    # Store result into json file
    with open(path_to_store, "w") as stream:
        json.dump(data, stream)