Source code for cgl.plugins.otio.tools.picker

import sys
import os

from PySide6 import QtCore, QtWidgets

from cgl.core.utils.general import load_json
from cgl.plugins.otio.tools import summary




class TreeItem(QtWidgets.QTreeWidgetItem):
    def __init__(self, parent, value):
        super(TreeItem, self).__init__(parent)
        self.value = value

        self.setFlags(self.flags() | QtCore.Qt.ItemIsUserCheckable)
        self.setText(0, value["sequence_name"])
        self.setText(1, value["shot_number"])
        self.setText(3, value.get("task", {}).get("msd", ""))

        # imported = value.get("task", {}).get("imported", False)
        # self.setCheckState(0, QtCore.Qt.Unchecked if imported else QtCore.Qt.Checked)
        self.setCheckState(0, QtCore.Qt.Unchecked)

        button = QtWidgets.QPushButton("Preview")

        button.clicked.connect(self.preview)
        parent.setItemWidget(self, 2, button)



    def is_checked(self):
        return self.checkState(0) == QtCore.Qt.Checked




    def preview(self):
        msd_path = self.value.get("task", {}).get("msd", None)
        root_path = self.value.get("root_path", None)
        if not msd_path or not root_path:
            QtWidgets.QMessageBox.critical(
                None, "Error!", "No msd_path or root_path", QtWidgets.QMessageBox.Ok
            )
            return

        full_msd_path = os.path.join(root_path, msd_path)
        if os.path.exists(full_msd_path):
            summary.show_summary_from_msd(full_msd_path, root_path)
        else:
            QtWidgets.QMessageBox.critical(
                None,
                "Error!",
                f"MSD path does not exist: {full_msd_path}",
                QtWidgets.QMessageBox.Ok,
            )






class StoryBoardPicker(QtWidgets.QDialog):
    def __init__(self, items, parent=None):
        super(StoryBoardPicker, self).__init__(parent)

        self.result = []

        self.tree = QtWidgets.QTreeWidget()
        self.tree.setHeaderLabels(["Sequence", "Shot", "", "Path"])

        self.tree.setSelectionMode(QtWidgets.QTreeWidget.ExtendedSelection)
        self.tree.setSelectionBehavior(QtWidgets.QTreeWidget.SelectRows)
        self.tree.customContextMenuRequested.connect(self.open_menu)
        self.tree.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)

        self.tree.setSortingEnabled(True)
        # for value in items:
        #     item = TreeItem(self.tree, value)

        self.button_box = QtWidgets.QDialogButtonBox(
            QtWidgets.QDialogButtonBox.Ok | QtWidgets.QDialogButtonBox.Cancel
        )
        self.button_box.accepted.connect(self.accept)
        self.button_box.rejected.connect(self.close)

        button_layout = QtWidgets.QHBoxLayout()
        button_layout.addStretch()
        button_layout.addWidget(self.button_box)

        layout = QtWidgets.QVBoxLayout()
        layout.addWidget(self.tree)
        layout.addLayout(button_layout)
        self.setLayout(layout)

        # Initialize
        self.resize(1000, 300)



    def accept(self):
        self.result = []
        it = QtWidgets.QTreeWidgetItemIterator(self.tree)
        while it.value():
            child = it.value()
            if child.is_checked():
                self.result.append(child.value)
            it += 1

        self.close()




    def open_menu(self, position):
        menu = QtWidgets.QMenu()
        check_action = QtWidgets.QAction("Check Selected")
        check_action.triggered.connect(lambda: self.set_selection(True))
        menu.addAction(check_action)
        uncheck_action = QtWidgets.QAction("Uncheck Selected")
        uncheck_action.triggered.connect(lambda: self.set_selection(False))
        menu.addAction(uncheck_action)
        menu.exec(self.tree.viewport().mapToGlobal(position))




    def set_selection(self, checked):
        it = QtWidgets.QTreeWidgetItemIterator(
            self.tree, QtWidgets.QTreeWidgetItemIterator.Selected
        )
        while it.value():
            item = it.value()
            item.setCheckState(0, QtCore.Qt.Checked if checked else QtCore.Qt.Unchecked)
            it += 1






def read_project_sequences(project_file_path, root_path):
    json_obj = load_json(project_file_path) or {}
    picker_items = []
    for sequence_name, shots in json_obj.get("data", {}).get("shots", {}).items():
        for shot_number, tasks in shots.items():
            edit_task = tasks.get("amc", {})
            if edit_task:
                item = {
                    "sequence_name": sequence_name,
                    "shot_number": shot_number,
                    "task": edit_task,
                    "root_path": root_path,
                }
                picker_items.append(item)

    return picker_items





def run_picker(project_file_path, root_path):
    picker_items = read_project_sequences(project_file_path, root_path)
    app = QtWidgets.QApplication.instance()
    if not app:
        app = QtWidgets.QApplication(sys.argv)

    if not picker_items:
        QtWidgets.QMessageBox.critical(
            None,
            "Error!",
            "Couldn't find any Editorial exports in Project msd file",
            QtWidgets.QMessageBox.Abort,
        )
        return []

    w = StoryBoardPicker(picker_items)
    w.setWindowTitle("Pick Animatic3D Sequences to import")
    w.exec()

    print(w.result)
    return w.result



if __name__ == "__main__":

    project_file = "C:\\CGLUMBERJACK\\COMPANIES\\premise\\render\\AlchemyTestB\\master\\project.msd"

    result = run_picker(project_file, "C:\\CGLUMBERJACK\\COMPANIES\\")

    # path_object.project