pymol-open-source/modules/pmg_qt/scene_bin_gui.py

from enum import IntEnum
from dataclasses import dataclass

from pymol.Qt import QtGui, QtWidgets
from pymol.Qt import QtCore
Qt = QtCore.Qt
QSI = QtGui.QStandardItem  # For brevity


class SceneDictIndex(IntEnum):
    QPIXMAP = 0
    MESSAGE = 1
    ACTIONS = 2


class SceneTableColumn(IntEnum):
    NAME = 0
    IMAGE = 1
    MESSAGE = 2
    ACTIONS = 3


@dataclass
class SceneWithTablePosition:
    name: str
    position: int


class ScenePanel(QtWidgets.QWidget):
    '''
    Scene Panel dialog for displaying and altering scenes in PyMOL.
    '''

    def __init__(self, parent):
        QtWidgets.QWidget.__init__(self)
        self.cmd = parent.cmd
        self.scene_list = self._update_scene_list()
        # EXAMPLE: scene_dict = {name:[QPixmap,Message,Actions]}
        self.scene_dict = {}

        self._build_table_elements(parent)
        self._populate_data()
        self._set_button_connections()

        self.resize(365, 700)

    def _build_table_elements(self, parent):
        '''
        Create the various elements in the table and add them to the
        proper layout.
        '''
        self.setWindowTitle('Scene Panel')
        layout = QtWidgets.QVBoxLayout(self)
        self.setLayout(layout)

        top_layout = QtWidgets.QGridLayout()
        layout.addLayout(top_layout)
        mid_layout = QtWidgets.QGridLayout()
        layout.addLayout(mid_layout)
        low_layout = QtWidgets.QGridLayout()
        layout.addLayout(low_layout)

        # Top Elements
        self.instructionLabel = QtWidgets.QLabel(self)
        self.instructionLabel.setText(
            'Double click selected thumbnail to \nload into Workspace.')
        top_layout.addWidget(self.instructionLabel, 0, 0)

        self.addSceneButton = QtWidgets.QPushButton(self)
        self.addSceneButton.setText('Add Scene')
        top_layout.addWidget(self.addSceneButton, 0, 1)

        # Mid Elements
        self.sceneTableWidget = QtWidgets.QTableWidget(self)
        mid_layout.addWidget(self.sceneTableWidget, 0, 0)
        self.sceneTableWidget.viewport().installEventFilter(self)

        self.sceneTableWidget.itemChanged.connect(self._item_changed)
        self.sceneTableWidget.selectionModel().selectionChanged.connect(
            self._selection_changed)
        self.sceneTableWidget.setSelectionBehavior(
            QtWidgets.QAbstractItemView.SelectionBehavior.SelectRows)
        self.sceneTableWidget.verticalHeader().setSectionsMovable(True)

        # Lower Buttom Elements
        self.deleteButton = QtWidgets.QPushButton(self)
        self.deleteButton.setText("Delete Scene")
        low_layout.addWidget(self.deleteButton, 0, 1)
        self.deleteButton.setEnabled(False)

        self.updateButton = QtWidgets.QPushButton(self)
        self.updateButton.setText("Update Scene")
        low_layout.addWidget(self.updateButton, 0, 0)
        self.updateButton.setEnabled(False)

    def _set_button_connections(self):
        self.addSceneButton.clicked.connect(self._add_scene)
        self.deleteButton.clicked.connect(self._delete_scene)
        self.updateButton.clicked.connect(self._update_scene)
        self.sceneTableWidget.doubleClicked.connect(self._show_selected_scene)

    def eventFilter(self, source, event):
        '''
        Event filter for capturing and processing events.
        '''
        focus_event = 24  # This will be replaced with the event itself.
        paint_event = 12
        if event.type() == paint_event:
            self._check_table_state()
        elif event.type() == focus_event:
            self._populate_data()

        return super().eventFilter(source, event)

    def _populate_data(self):
        '''
        Main loop for populating the table and class data structures.
        '''
        # Clear rows
        self.sceneTableWidget.setRowCount(0)
        self.scene_list = self._update_scene_list()
        self._update_scene_dict()

        self._format_condensed_widget()

        for scene_name in self.scene_list:
            # Add a row for the scene
            row_number = self.sceneTableWidget.rowCount()
            self.sceneTableWidget.insertRow(row_number)
            self.sceneTableWidget.setRowHeight(row_number, 100)
            # Get data from scene_dict
            scene_image_label = self._get_scene_image_label(
                self.scene_dict[scene_name][SceneDictIndex.QPIXMAP])
            scene_message = self.scene_dict[scene_name][SceneDictIndex.MESSAGE]
            scene_actions = self.scene_dict[scene_name][SceneDictIndex.ACTIONS]

            # Set items within the new row
            self.sceneTableWidget.setItem(
                row_number, SceneTableColumn.NAME,
                QtWidgets.QTableWidgetItem(scene_name))
            self.sceneTableWidget.setCellWidget(
                row_number, SceneTableColumn.IMAGE, scene_image_label)
            self.sceneTableWidget.setItem(
                row_number, SceneTableColumn.MESSAGE,
                QtWidgets.QTableWidgetItem(scene_message))
            self.sceneTableWidget.setItem(
                row_number, SceneTableColumn.ACTIONS,
                QtWidgets.QTableWidgetItem(scene_actions))

        self._format_table()

        self.sceneTableWidget.selectionModel().clearSelection()

    def _update_scene_dict(self):
        '''
        Updates the class' scene_dict to reflect the current state of
        the class' scene_list. Currently, place holder value are being used
        for the message and action but these will be replaced by functions
        responsible for fetching these values for a given scene.

        Note: This is also currently being used to make sure that scenes
        are being properly removed from the dictionary if they no longer
        appear in the scene_list. There may be a better way of doing this or
        this may be handled in separate step in the future.
        '''
        for scene_position in range(len(self.scene_list)):
            scene_name = self.scene_list[scene_position]
            if scene_name not in self.scene_dict:
                scene_pix_map = self._get_scene_png(scene_name)
                scene_message = 'This is a base message'
                scene_actions = 'Rock, zoom, something'
                self.scene_dict[scene_name] = [
                    scene_pix_map, scene_message,
                    scene_actions, scene_position]

        remove_scenes = [
            scene_name for scene_name in self.scene_dict
            if scene_name not in self.scene_list]

        for scene_name in remove_scenes:
            self.scene_dict.pop(scene_name)

    def _get_scene_png(self, scene_name):
        '''
        Returns a QPixmap object that can be assigned to a label
        in order to display the scene thumbnail.
        This will be obsolete and removed in the future.
        '''
        scene_pix_map = QtGui.QPixmap()
        png_buf = self.cmd.get_scene_thumbnail(scene_name)
        scene_pix_map.loadFromData(png_buf, "PNG")
        return scene_pix_map

    def _get_scene_image_label(self, scene_pix_map):
        '''
        Takes in the QPixmap from get_scene_png and uses it to create a label
        displaying the thumbnail.
        In the future, this function will generate its
        own QPixmap object from the PNG buffer stored with the scene.
        '''
        scene_image_label = QtWidgets.QLabel(self)
        scene_image_label.setPixmap(scene_pix_map)
        return scene_image_label

    def _format_condensed_widget(self):
        '''
        Formats the widget to look and size properly. There will be a
        matching function for the expanded view of the widget.
        '''
        self.sceneTableWidget.setMinimumWidth(275)
        self.setMinimumWidth(375)
        self.resize(300, self.height())
        self.sceneTableWidget.setColumnCount(2)
        self.sceneTableWidget.setHorizontalHeaderLabels(
            ['Name', 'Scene Preview'])
        self.sceneTableWidget.resize(260, 500)

    def _format_table(self):
        '''
        Formats the table to look properly. Currently, vertical headers
        are being hidden but they will be used for reordering.
        '''
        hh = self.sceneTableWidget.horizontalHeader()
        hh.setStretchLastSection(True)
        self.sceneTableWidget.setFocus()
        self.sceneTableWidget.hide()
        self.sceneTableWidget.resizeColumnsToContents()
        self._set_vertical_headers()
        self.sceneTableWidget.show()

    def _set_vertical_headers(self):
        # Get row count
        row_count = self.sceneTableWidget.rowCount()

        # Iterate in range, setting item
        for ind in range(row_count):
            # Create  arrows item
            arrows = QtWidgets.QTableWidgetItem('\u2195')
            arrows_font = arrows.font()
            arrows_font.setPointSize(20)
            arrows.setFont(arrows_font)
            self.sceneTableWidget.setVerticalHeaderItem(ind, arrows)

    def _update_scene_list(self):
        '''
        Returns an updated list of the current scenes. Currently, this
        is very basic but some handling will be done here in the future
        or this will likely be removed.
        '''
        list_scenes = self.cmd.get_scene_list()
        return list_scenes

    def _add_scene(self):
        '''
        Calls cmd.scene with append and repopulates the table.
        Table is scrolled to the bottom to show new scene.
        '''
        self.cmd.scene('new', 'append', quiet=0)
        self._populate_data()
        self.sceneTableWidget.scrollToBottom()

    def _update_scene(self):
        '''
        Calls cmd.scene with update, replacing the selected scene with
        new scene. Image is not currently being updated but it should.
        Button only active with selection.
        '''
        selection = self.sceneTableWidget.selectionModel().selectedIndexes()
        if selection:
            name = selection[0].data()
            self.cmd.scene(name, 'update')
            self.scene_dict[name][SceneDictIndex.QPIXMAP] = self._get_scene_png(
                name)
            self._update_table()

    def _delete_scene(self):
        '''
        Calls cmd.scene with clear, removing the selected scene.
        Button only active with selection.
        '''
        selection = self.sceneTableWidget.selectionModel().selectedIndexes()
        name = selection[0].data()
        self.cmd.scene(name, 'clear')
        try:
            self.scene_dict.pop(name)
            self.scene_list.remove(name)
        except Exception as e:
            print("Item not found")
            print(e)
        self._update_table()

    def _show_selected_scene(self):
        '''
        Calls cmd.scene with recall, recalling scene to workspace.
        '''
        selection = self.sceneTableWidget.selectionModel().selectedIndexes()
        if selection:
            name = selection[0].data()
            self.cmd.scene(name, 'recall')

    def _update_table(self):
        '''
        Currently pointless but scrolling/formatting will happen
        here to make updating the table more seemless after
        intially populating.
        '''
        self._populate_data()

    def _compare_scene_lists(self, old_list, new_list):
        '''
        Compares current and new list of scenes.
        Returns list of elements that differ ([old, new])
        '''
        diff_list = []
        for i in range(min(len(old_list), len(new_list))):
            if new_list[i] != old_list[i]:
                diff_list.append([old_list[i], new_list[i]])
        return diff_list

    def _get_table_scene_list(self):
        '''
        Returns a sorted list of scene names based on the current
        state of the table. This will be a critical part of reordering,
        but also used for detecting renaming.
        '''
        scene_coor_list = []
        for row_num in range(self.sceneTableWidget.rowCount()):
            name = self.sceneTableWidget.model().data(
                self.sceneTableWidget.model().index(row_num, 0))
            position = int(
                self.sceneTableWidget.rowViewportPosition(row_num)/100)

            scene_coor_list.append(SceneWithTablePosition(name, position))

        scene_coor_list.sort(key=lambda scene: scene.position)

        return [scene_with_pos.name for scene_with_pos in scene_coor_list]

    def _check_table_state(self):
        '''
        Checks the current state of the table for reordering.
        '''
        scene_ordered_list = self._get_table_scene_list()
        self.scene_list = self.cmd.get_scene_list()
        diff_list = self._compare_scene_lists(
            self.scene_list, scene_ordered_list)
        if len(diff_list) > 1:
            # This is the reordering case since multiple have changed
            self._reorder_scenes(scene_ordered_list)

    def _selection_changed(self):
        '''
        Enables buttons when selection is made.
        '''
        item_selected = bool(
            self.sceneTableWidget.selectionModel().selectedIndexes())
        self.deleteButton.setEnabled(item_selected)
        self.updateButton.setEnabled(item_selected)

    def _rename_scene(self, item):
        '''
        Takes in an item from _item_changed and checks conditions before
        calling scene with 'rename'.
        '''
        if ' ' in item.text():
            print("Scene names with spaces are not supported")
            self._update_table()
        elif not item.text():
            print("Blank scene names are not allowed")
            self._update_table()
        else:
            diff_list = self._compare_scene_lists(
                self.cmd.get_scene_list(),
                self._get_table_scene_list())
            if diff_list:
                for old_scene, new_scene in diff_list:
                    self.cmd.scene(old_scene, 'rename', new_key=new_scene)

    def _reorder_scenes(self, scene_ordered_list):
        '''
        Calls cmd.scene_order which requires a string of scene names.
        This string is created using the ordered scene list generated
        by _check_table_state.
        '''
        scene_order_string = ' '.join(
            scene for scene in scene_ordered_list if scene in self.scene_list)
        self.cmd.scene_order(scene_order_string)
        self.scene_list = self._update_scene_list()

    def _item_changed(self, item):
        """
        Called every time an item in the table is changed.
        @param item: The item which has changed
        @type  item: QStandardItem
        """
        if item.column() == SceneTableColumn.NAME:
            self._rename_scene(item)