random_repace.py

"""
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
"""

import platform
import random
from functools import partial

import maya.OpenMayaUI as omui
import pymel.core as pm
from PySide2 import QtCore
from PySide2 import QtWidgets
from shiboken2 import wrapInstance


__version__ = '1.1.0'

WINDOW_NAME = "Random Replace - v{}".format(__version__)

# Python 3 compatibility
if platform.python_version_tuple()[0] == '3':
    long = int


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)


def close_existing_windows():
    """
    Close any existing instances of the window
    """
    for child_window in maya_main_window().children():
        if hasattr(child_window, 'windowTitle'):
            if child_window.windowTitle() == WINDOW_NAME:
                child_window.close()
                child_window.deleteLater()


class RandomReplace(QtWidgets.QDialog):
    """
    Main UI Class for the importer
    """

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

        self.setWindowTitle(WINDOW_NAME)
        self.setMinimumWidth(300)
        self.setWindowFlags(self.windowFlags() ^ QtCore.Qt.WindowContextHelpButtonHint)

        self.create_widgets()
        self.create_layouts()
        self.create_connections()
        self.source_objects = []
        self.target_objects = []

    def create_widgets(self):
        self.mode_combo_box = QtWidgets.QComboBox()
        self.mode_combo_box.addItems(['Replace', 'Add', 'Hide'])
        self.use_instance_cb = QtWidgets.QCheckBox('Use Instances (Default duplicates)')
        self.source_line_edit = QtWidgets.QLineEdit()
        self.source_btn = QtWidgets.QPushButton('Source')
        self.target_line_edit = QtWidgets.QLineEdit()
        self.target_btn = QtWidgets.QPushButton('Target')
        self.replace_btn = QtWidgets.QPushButton('Replace')

        self.source_btn.setFixedHeight(20)
        self.source_line_edit.setFixedHeight(20)
        self.source_line_edit.setEnabled(False)
        self.target_btn.setFixedHeight(20)
        self.target_line_edit.setFixedHeight(20)
        self.target_line_edit.setEnabled(False)
        self.replace_btn.setMinimumHeight(40)

    def create_layouts(self):
        settings_layout = QtWidgets.QVBoxLayout()
        settings_layout.addWidget(self.mode_combo_box)
        settings_layout.addWidget(self.use_instance_cb)

        form_layout = QtWidgets.QFormLayout()
        form_layout.addRow(self.source_btn, self.source_line_edit)
        form_layout.addRow(self.target_btn, self.target_line_edit)

        button_layout = QtWidgets.QHBoxLayout()
        button_layout.addWidget(self.replace_btn)

        main_layout = QtWidgets.QVBoxLayout(self)
        main_layout.setContentsMargins(4, 4, 4, 4)
        main_layout.addLayout(settings_layout)
        main_layout.addLayout(form_layout)
        main_layout.addLayout(button_layout)

    def create_connections(self):
        self.source_btn.clicked.connect(partial(self.set_line_edit, self.source_line_edit))
        self.target_btn.clicked.connect(partial(self.set_line_edit, self.target_line_edit))
        self.replace_btn.clicked.connect(self.replace_objects)

    def set_line_edit(self, line_edit):
        """
        Capture the selection and set it to both internal variables and the line edit
        :param line_edit: The target line edit.
        """
        # Setting up the dialog filters to only accept what is expected in that field to prevent user error.
        user_selection = pm.ls(selection=True)
        if line_edit == self.source_line_edit:
            if len(user_selection) > 0:
                self.source_objects = user_selection
                text = str([x.name() for x in self.source_objects])
            else:
                self.source_objects = []
                text = ''
        else:
            if len(user_selection) > 0:
                self.target_objects = user_selection
                text = str([x.name() for x in self.target_objects])
            else:
                self.target_objects = []
                text = ''

        line_edit.setText(text)

    def replace_objects(self):
        """
        Main function to generate the camera and image plane from UI.
        """
        print(self.mode_combo_box.currentText())
        print(self.use_instance_cb.isChecked())

        if not self.source_objects:
            pm.warning('Please select objects to replace with, then click the "Source" to load them')

        elif not self.target_objects:
            pm.warning('Please select objects to replace, then click the "Target" to load them')

        else:
            mode = self.mode_combo_box.currentText()
            use_instance = self.use_instance_cb.isEnabled()

            for target in self.target_objects:
                source = random.choice(self.source_objects)

                new_source = pm.instance(source) if use_instance else pm.duplicate(source)
                pm.matchTransform(new_source, target)

                if mode == 'Replace':
                    pm.delete(target)
                if mode == 'Hide':
                    pm.hide(target)

            print('Done!')


if __name__ == '__main__':
    close_existing_windows()
    random_replace_dialog = RandomReplace()
    random_replace_dialog.show()