Viewport Overlay

# -*- coding: utf-8 -*-
"""
Author: Virgil
Date: 2020/6/20 18:13
"""

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 ViewportOverlay(QtWidgets.QWidget):
    def __init__(self):
        super(ViewportOverlay, self).__init__(None)

        self.update_active_viewport()
        self.set_context_menu_enabled(False)

    def update_active_viewport(self):
        active_viewport_ptr = omui.M3dView.active3dView().widget()
        self.viewport_widget = wrapInstance(long(active_viewport_ptr), QtWidgets.QWidget)

        self.setParent(self.viewport_widget)
        self.setGeometry(self.viewport_widget.geometry())

    def set_context_menu_enabled(self, enabled):
        self.context_menu_enabled = enabled

    def paintEvent(self, event):
        if self.viewport_widget:
            painter = QtGui.QPainter(self)
            painter.fillRect(100, 100, 200, 200, QtGui.QColor(QtCore.Qt.green))

if __name__ == "__main__":
    try:
        ViewportOverlayDialog.delete_overlays()
    except:
        pass

class ViewportOverlayDialog(QtWidgets.QDialog):
    WINDOW_TITLE = "Viewport Overlay"

    viewport_overlay = None

    @classmethod
    def delete_overlays(cls):
        if ViewportOverlayDialog.viewport_overlay:
            ViewportOverlayDialog.viewport_overlay.setParent(None)
            ViewportOverlayDialog.viewport_overlay.deleteLater()
            ViewportOverlayDialog.viewport_overlay = None

    def __init__(self, parent=maya_main_window()):
        super(ViewportOverlayDialog, self).__init__(parent)

        self.setWindowTitle(self.WINDOW_TITLE)
        self.setMinimumSize(280, 160)

        self.create_widgets()
        self.create_layout()
        self.create_connections()

    def create_widgets(self):
        self.overlay_visible_cb = QtWidgets.QCheckBox("Show Overlay")

        self.context_menu_cb = QtWidgets.QCheckBox("Context Menu Enabled")
        self.context_menu_cb.setChecked(True)

        self.close_bttn = QtWidgets.QPushButton("Close")

    def create_layout(self):

        overlay_layout = QtWidgets.QVBoxLayout()
        overlay_layout.addWidget(self.overlay_visible_cb)
        overlay_layout.addWidget(self.context_menu_cb)

        options_grp = QtWidgets.QGroupBox("Overlay Options")
        options_grp.setLayout(overlay_layout)

        bttn_layout = QtWidgets.QHBoxLayout()
        bttn_layout.addStretch()
        bttn_layout.addWidget(self.close_bttn)

        main_layout = QtWidgets.QVBoxLayout(self)
        main_layout.setContentsMargins(5, 5, 5, 5)
        main_layout.addWidget(options_grp)
        main_layout.addStretch()
        main_layout.addLayout(bttn_layout)

    def create_connections(self):
        self.overlay_visible_cb.toggled.connect(self.set_overlay_visible)
        self.close_bttn.clicked.connect(self.close)

    def set_overlay_visible(self, visible):
        if visible:
            if not ViewportOverlayDialog.viewport_overlay:
                ViewportOverlayDialog.viewport_overlay = ViewportOverlay()
                ViewportOverlayDialog.viewport_overlay.set_context_menu_enabled(self.context_menu_cb.isChecked())

                self.context_menu_cb.toggled.connect(ViewportOverlayDialog.viewport_overlay.set_context_menu_enabled)

        if ViewportOverlayDialog.viewport_overlay:
            ViewportOverlayDialog.viewport_overlay.setVisible(visible)

        self.overlay_visible_cb.setChecked(visible)
    
def showWindow():
    global dialog
    try:
        dialog.close()
        dialog.deleteLater()
    except:
        pass

    dialog = ViewportOverlayDialog()
    dialog.show()

if __name__ == "__main__":
    try:
        dialog.close()
        dialog.deleteLater()
    except:
        pass

    dialog = ViewportOverlayDialog()
    dialog.show()