PimpMySpreadsheet

class CustomSpreadsheetColumns(QObject):
    """
    A class defining custom columns for Hiero's spreadsheet view. This has a similar, but
    slightly simplified, interface to the QAbstractItemModel and QItemDelegate classes.
  """
    global gStatusTags
    global gArtistList

    # Ideally, we'd set this list on a Per Item basis, but this is expensive for a large mixed selection
    standardColourSpaces = [
        "linear", "sRGB", "rec709", "Cineon", "Gamma1.8", "Gamma2.2",
        "Panalog", "REDLog", "ViperLog"
    ]
    arriColourSpaces = [
        "Video - Rec709", "LogC - Camera Native", "Video - P3", "ACES",
        "LogC - Film", "LogC - Wide Gamut"
    ]
    r3dColourSpaces = [
        "Linear", "Rec709", "REDspace", "REDlog", "PDlog685", "PDlog985",
        "CustomPDlog", "REDgamma", "SRGB", "REDlogFilm", "REDgamma2",
        "REDgamma3"
    ]
    gColourSpaces = standardColourSpaces + arriColourSpaces + r3dColourSpaces

    currentView = hiero.ui.activeView()

    # This is the list of Columns available
    gCustomColumnList = [
        {
            "name": "Tags",
            "cellType": "readonly"
        },
        {
            "name": "Colourspace",
            "cellType": "dropdown"
        },
        {
            "name": "Notes",
            "cellType": "readonly"
        },
        {
            "name": "FileType",
            "cellType": "readonly"
        },
        {
            "name": "Shot Status",
            "cellType": "dropdown"
        },
        {
            "name": "Thumbnail",
            "cellType": "readonly"
        },
        {
            "name": "MediaType",
            "cellType": "readonly"
        },
        {
            "name": "Width",
            "cellType": "readonly"
        },
        {
            "name": "Height",
            "cellType": "readonly"
        },
        {
            "name": "Pixel Aspect",
            "cellType": "readonly"
        },
        {
            "name": "Artist",
            "cellType": "dropdown"
        },
        {
            "name": "Department",
            "cellType": "readonly"
        },
    ]

    def numColumns(self):
        """
      Return the number of custom columns in the spreadsheet view
    """
        return len(self.gCustomColumnList)

    def columnName(self, column):
        """
      Return the name of a custom column
    """
        return self.gCustomColumnList[column]["name"]

    def getTagsString(self, item):
        """
      Convenience method for returning all the Notes in a Tag as a string
    """
        tagNames = []
        tags = item.tags()
        for tag in tags:
            tagNames += [tag.name()]
        tagNameString = ','.join(tagNames)
        return tagNameString

    def getNotes(self, item):
        """
      Convenience method for returning all the Notes in a Tag as a string
    """
        notes = ""
        tags = item.tags()
        for tag in tags:
            note = tag.note()
            if len(note) > 0:
                notes += tag.note() + ', '
        return notes[:-2]

    def getData(self, row, column, item):
        """
      Return the data in a cell
    """
        currentColumn = self.gCustomColumnList[column]
        if currentColumn["name"] == "Tags":
            return self.getTagsString(item)

        if currentColumn["name"] == "Colourspace":
            try:
                colTransform = item.sourceMediaColourTransform()
            except:
                colTransform = "--"
            return colTransform

        if currentColumn["name"] == "Notes":
            try:
                note = self.getNotes(item)
            except:
                note = ""
            return note

        if currentColumn["name"] == "FileType":
            fileType = "--"
            M = item.source().mediaSource().metadata()
            if M.hasKey("foundry.source.type"):
                fileType = M.value("foundry.source.type")
            elif M.hasKey("media.input.filereader"):
                fileType = M.value("media.input.filereader")
            return fileType

        if currentColumn["name"] == "Shot Status":
            status = item.status()
            if not status:
                status = "--"
            return str(status)

        if currentColumn["name"] == "MediaType":
            M = item.mediaType()
            return str(M).split("MediaType")[-1].replace(".k", "")

        if currentColumn["name"] == "Thumbnail":
            return str(item.eventNumber())

        if currentColumn["name"] == "Width":
            return str(item.source().format().width())

        if currentColumn["name"] == "Height":
            return str(item.source().format().height())

        if currentColumn["name"] == "Pixel Aspect":
            return str(item.source().format().pixelAspect())

        if currentColumn["name"] == "Artist":
            if item.artist():
                name = item.artist()["artistName"]
                return name
            else:
                return "--"

        if currentColumn["name"] == "Department":
            if item.artist():
                dep = item.artist()["artistDepartment"]
                return dep
            else:
                return "--"

        return ""

    def setData(self, row, column, item, data):
        """
      Set the data in a cell - unused in this example
    """

        return None

    def getTooltip(self, row, column, item):
        """
      Return the tooltip for a cell
    """
        currentColumn = self.gCustomColumnList[column]
        if currentColumn["name"] == "Tags":
            return str([item.name() for item in item.tags()])

        if currentColumn["name"] == "Notes":
            return str(self.getNotes(item))
        return ""

    def getFont(self, row, column, item):
        """
      Return the tooltip for a cell
    """
        return None

    def getBackground(self, row, column, item):
        """
      Return the background colour for a cell
    """
        if not item.source().mediaSource().isMediaPresent():
            return QColor(80, 20, 20)
        return None

    def getForeground(self, row, column, item):
        """
      Return the text colour for a cell
    """
        #if column == 1:
        #  return QColor(255, 64, 64)
        return None

    def getIcon(self, row, column, item):
        """
      Return the icon for a cell
    """
        currentColumn = self.gCustomColumnList[column]
        if currentColumn["name"] == "Colourspace":
            return QIcon("icons:LUT.png")

        if currentColumn["name"] == "Shot Status":
            status = item.status()
            if status:
                return QIcon(gStatusTags[status])

        if currentColumn["name"] == "MediaType":
            mediaType = item.mediaType()
            if mediaType == hiero.core.TrackItem.kVideo:
                return QIcon("icons:VideoOnly.png")
            elif mediaType == hiero.core.TrackItem.kAudio:
                return QIcon("icons:AudioOnly.png")

        if currentColumn["name"] == "Artist":
            try:
                return QIcon(item.artist()["artistIcon"])
            except:
                return None
        return None

    def getSizeHint(self, row, column, item):
        """
      Return the size hint for a cell
    """
        currentColumnName = self.gCustomColumnList[column]["name"]

        if currentColumnName == "Thumbnail":
            return QSize(90, 50)

        return QSize(50, 50)

    def paintCell(self, row, column, item, painter, option):
        """
      Paint a custom cell. Return True if the cell was painted, or False to continue
      with the default cell painting.
    """
        currentColumn = self.gCustomColumnList[column]
        if currentColumn["name"] == "Tags":
            if option.state & QStyle.State_Selected:
                painter.fillRect(option.rect, option.palette.highlight())
            iconSize = 20
            r = QRect(option.rect.x(),
                      option.rect.y() + (option.rect.height() - iconSize) / 2,
                      iconSize, iconSize)
            tags = item.tags()
            if len(tags) > 0:
                painter.save()
                painter.setClipRect(option.rect)
                for tag in item.tags():
                    M = tag.metadata()
                    if not (M.hasKey("tag.status")
                            or M.hasKey("tag.artistID")):
                        QIcon(tag.icon()).paint(painter, r, Qt.AlignLeft)
                        r.translate(r.width() + 2, 0)
                painter.restore()
                return True

        if currentColumn["name"] == "Thumbnail":
            imageView = None
            pen = QPen()
            r = QRect(option.rect.x() + 2, (option.rect.y() +
                                            (option.rect.height() - 46) / 2),
                      85, 46)
            if not item.source().mediaSource().isMediaPresent():
                imageView = QImage("icons:Offline.png")
                pen.setColor(QColor(Qt.red))

            if item.mediaType() == hiero.core.TrackItem.MediaType.kAudio:
                imageView = QImage("icons:AudioOnly.png")
                #pen.setColor(QColor(Qt.green))
                painter.fillRect(r, QColor(45, 59, 45))

            if option.state & QStyle.State_Selected:
                painter.fillRect(option.rect, option.palette.highlight())

            tags = item.tags()
            painter.save()
            painter.setClipRect(option.rect)

            if not imageView:
                try:
                    imageView = item.thumbnail(item.sourceIn())
                    pen.setColor(QColor(20, 20, 20))
                # If we're here, we probably have a TC error, no thumbnail, so get it from the source Clip...
                except:
                    pen.setColor(QColor(Qt.red))

            if not imageView:
                try:
                    imageView = item.source().thumbnail()
                    pen.setColor(QColor(Qt.yellow))
                except:
                    imageView = QImage("icons:Offline.png")
                    pen.setColor(QColor(Qt.red))

            QIcon(QPixmap.fromImage(imageView)).paint(painter, r,
                                                      Qt.AlignCenter)
            painter.setPen(pen)
            painter.drawRoundedRect(r, 1, 1)
            painter.restore()
            return True

        return False

    def createEditor(self, row, column, item, view):
        """
      Create an editing widget for a custom cell
    """
        self.currentView = view

        currentColumn = self.gCustomColumnList[column]
        if currentColumn["cellType"] == "readonly":
            cle = QLabel()
            cle.setEnabled(False)
            cle.setVisible(False)
            return cle

        if currentColumn["name"] == "Colourspace":
            cb = QComboBox()
            for colourspace in self.gColourSpaces:
                cb.addItem(colourspace)
            cb.currentIndexChanged.connect(self.colourspaceChanged)
            return cb

        if currentColumn["name"] == "Shot Status":
            cb = QComboBox()
            cb.addItem("")
            for key in gStatusTags.keys():
                cb.addItem(QIcon(gStatusTags[key]), key)
            cb.addItem("--")
            cb.currentIndexChanged.connect(self.statusChanged)

            return cb

        if currentColumn["name"] == "Artist":
            cb = QComboBox()
            cb.addItem("")
            for artist in gArtistList:
                cb.addItem(artist["artistName"])
            cb.addItem("--")
            cb.currentIndexChanged.connect(self.artistNameChanged)
            return cb
        return None

    def setModelData(self, row, column, item, editor):
        return False

    def dropMimeData(self, row, column, item, data, items):
        """
      Handle a drag and drop operation - adds a Dragged Tag to the shot
    """
        for thing in items:
            if isinstance(thing, hiero.core.Tag):
                item.addTag(thing)
        return None

    def colourspaceChanged(self, index):
        """
      This method is called when Colourspace widget changes index.
    """
        index = self.sender().currentIndex()
        colourspace = self.gColourSpaces[index]
        selection = self.currentView.selection()
        project = selection[0].project()
        with project.beginUndo("Set Colourspace"):
            items = [
                item for item in selection
                if (item.mediaType() == hiero.core.TrackItem.MediaType.kVideo)
            ]
            for trackItem in items:
                trackItem.setSourceMediaColourTransform(colourspace)

    def statusChanged(self, arg):
        """
      This method is called when Shot Status widget changes index.
    """
        view = hiero.ui.activeView()
        selection = view.selection()
        status = self.sender().currentText()
        project = selection[0].project()
        with project.beginUndo("Set Status"):
            # A string of "--" characters denotes clear the status
            if status != "--":
                for trackItem in selection:
                    trackItem.setStatus(status)
            else:
                for trackItem in selection:
                    tTags = trackItem.tags()
                    for tag in tTags:
                        if tag.metadata().hasKey("tag.status"):
                            trackItem.removeTag(tag)
                            break

    def artistNameChanged(self, arg):
        """
      This method is called when Artist widget changes index.
    """
        view = hiero.ui.activeView()
        selection = view.selection()
        name = self.sender().currentText()
        project = selection[0].project()
        with project.beginUndo("Assign Artist"):
            # A string of "--" denotes clear the assignee...
            if name != "--":
                for trackItem in selection:
                    trackItem.setArtistByName(name)
            else:
                for trackItem in selection:
                    tTags = trackItem.tags()
                    for tag in tTags:
                        if tag.metadata().hasKey("tag.artistID"):
                            trackItem.removeTag(tag)
                            break