Timeline Overlay

# -*- coding: utf-8 -*-
"""
Author: Virgil
Date: 2020/6/20 15:17
"""

from PySide2 import QtCore
from PySide2 import QtGui
from PySide2 import QtWidgets
from shiboken2 import wrapInstance

import maya.OpenMayaUI as omui
import maya.cmds as cmds
import maya.mel as mel

def maya_main_window():
    """
    Return the Maya main window widget as a Python object
    """
    main_window_ptr = omui.MQtUtil.mainWindow()
    return wrapInstance(long(main_window_ptr), QtWidgets.QWidget)

class TimelineOverlay(QtWidgets.QWidget):

    KEYFRAME_COLOR = QtGui.QColor(QtCore.Qt.green)

    def __init__(self):

        self.time_control = mel.eval("$tempVar = $gPlayBackSlider")
        time_control_ptr = omui.MQtUtil.findControl(self.time_control)
        time_control_widget = wrapInstance(long(time_control_ptr), QtWidgets.QWidget)

        super(TimelineOverlay, self).__init__(time_control_widget)

        self.frame_times = [1, 5, 8, 10, 19, 40, 59, 60, 61, 93]

        self.set_context_menu_enabled(False)

    def add_frame(self):
        print("TODO: Add Frame")
        current_time = cmds.currentTime(query=True)
        if current_time not in self.frame_times:
            self.frame_times.append(current_time)
            self.update()

    def add_frames(self):
        print("TODO: Add Frames")
        selected_range = self.get_selected_range()
        for frame_time in range(selected_range[0], selected_range[1]):
            if frame_time not in self.frame_times:
                self.frame_times.append