vis_attr_maker.py

"""
VIS ATTR MAKER SCRIPT

# Open UI - Add this code to a button in Maya
---

import vis_attr_maker
ui = vis_attr_maker.VisSwitchMakerUI()
ui.show()

---

# Build Vis switch from exported File

---
import vis_attr_maker
path = r"<PATH TO YOUR EXPORTED CONFIG FILE>"
vis_attr_maker.create_from_file(path)
---

"""
import sys
import json

from PySide2 import QtCore
from PySide2 import QtUiTools
from PySide2 import QtWidgets
from shiboken2 import wrapInstance

import pymel.core as pm
import maya.OpenMayaUI as omui

__author__ = 'Justin Pedersen'
__version__ = 'v0.0.2'


def maya_main_window():
    """
    Return the Maya main window widget as a Python object
    """
    main_window_ptr = omui.MQtUtil.mainWindow()
    if sys.version_info.major >= 3:
        return wrapInstance(int(main_window_ptr), QtWidgets.QWidget)
    else:
        return wrapInstance(long(main_window_ptr), QtWidgets.QWidget)


def create_from_file(file_path):
    with open(file_path, 'r') as file:
        config_dict = json.load(file)

    # Validate the necessary keys exist in the dictionary
    necessary_keys = ['host_ctrl', 'attr_name', 'on_ctrls', 'off_ctrls']
    if not all(key in config_dict for key in necessary_keys):
        pm.warning("Invalid config file.")
        return

    connect_vis_attrs_from_dict(config_dict)


def connect_output(ctrl, output_attr):
    node = pm.PyNode(ctrl)
    # Connect the shape nodes
    shapes = pm.listRelatives(node, shapes=True)

    # If we are given a transform node with no shapes or a transform node with mesh objects
    if not shapes or all([s for s in shapes if pm.nodeType(s) == 'mesh']):
        pm.connectAttr(output_attr, node.visibility, force=True)

    # Plug into the vis of each shape node
    else:
        for shape in pm.listRelatives(node, shapes=True):
            pm.connectAttr(output_attr, shape.visibility, force=True)


def add_enum_attr(ctrl, attr_name):
    """
    """
    # Make the ctrl a PyBode if not already
    if not isinstance(ctrl, pm.PyNode):
        ctrl = pm.PyNode(ctrl)

    # If the attr doesn't exist, then add it
    if not ctrl.hasAttr(attr_name):
        ctrl.addAttr(attr_name, at='enum', en="On:Off", keyable=True)

    return ctrl.attr(attr_name)


def connect_vis_attrs_from_dict(data):
    host_ctrl = data['host_ctrl']
    attr_name = data['attr_name']

    # Add the host attr
    vis_attr = add_enum_attr(host_ctrl, attr_name)

    for ctrl in data['on_ctrls']:
        try:
            connect_output(ctrl, vis_attr)

        except Exception as e:
            pm.warning(f'Ctrl {ctrl} skipped with error: {e}')

    # create a not node to invert the output
    not_node = pm.createNode('not')
    pm.connectAttr(vis_attr, not_node.input, force=True)

    for ctrl in data['off_ctrls']:
        try:
            connect_output(ctrl, not_node.output)

        except Exception as e:
            pm.warning(f'Ctrl {ctrl} skipped with error: {e}')


# noinspection PyAttributeOutsideInit
class VisSwitchMakerUI(QtWidgets.QDialog):

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

        self.setWindowTitle("Vis Switch Maker")

        self._host_ui_ctrl_node = None
        self._on_ctrls = []
        self._off_ctrls = []

        self.create_widgets()
        self.create_layouts()
        self.create_connections()

    def create_widgets(self):
        self.host_ui_ctrl_label = QtWidgets.QLabel("Host UI Ctrl:")
        self.host_ui_ctrl_line_edit = QtWidgets.QLineEdit()
        self.set_host_ui_ctrl_button = QtWidgets.QPushButton('<<<')

        self.attr_name_label = QtWidgets.QLabel("Attr Name:")
        self.attr_line_edit = QtWidgets.QLineEdit()

        self.enum_options_combo = QtWidgets.QComboBox()
        self.enum_options_combo.addItems(['On', 'Off'])

        self.ctrls_group_box = QtWidgets.QGroupBox("Ctrls that will be visible on the above attr")
        self.ctrls_list = QtWidgets.QListWidget()

        self.ctrls_add_button = QtWidgets.QPushButton('Add')
        self.ctrls_remove_button = QtWidgets.QPushButton('Remove')

        self.export_config_button = QtWidgets.QPushButton('Export Config')
        self.import_config_button = QtWidgets.QPushButton('Import Config')

        self.create_button = QtWidgets.QPushButton('Create')

    def create_layouts(self):
        main_layout = QtWidgets.QVBoxLayout()
        self.setLayout(main_layout)

        # Attr layout
        host_ui_ctrl_layout = QtWidgets.QHBoxLayout()
        host_ui_ctrl_layout.addWidget(self.host_ui_ctrl_line_edit)
        host_ui_ctrl_layout.addWidget(self.set_host_ui_ctrl_button)

        host_ui_form_layout = QtWidgets.QFormLayout()
        host_ui_form_layout.addRow(self.host_ui_ctrl_label, host_ui_ctrl_layout)
        host_ui_form_layout.addRow(self.attr_name_label, self.attr_line_edit)

        main_layout.addLayout(host_ui_form_layout)

        # Ctrls layout
        add_remove_button_layout = QtWidgets.QHBoxLayout()
        add_remove_button_layout.addWidget(self.ctrls_add_button)
        add_remove_button_layout.addWidget(self.ctrls_remove_button)

        group_box_layout = QtWidgets.QVBoxLayout()
        group_box_layout.addWidget(self.ctrls_list)
        group_box_layout.addLayout(add_remove_button_layout)
        self.ctrls_group_box.setLayout(group_box_layout)

        main_ctrls_layout = QtWidgets.QVBoxLayout()
        main_ctrls_layout.addWidget(self.enum_options_combo)
        main_ctrls_layout.addWidget(self.ctrls_group_box)

        main_layout.addLayout(main_ctrls_layout)
        main_layout.addStretch()

        export_button_layout = QtWidgets.QHBoxLayout()
        export_button_layout.addWidget(self.export_config_button)
        export_button_layout.addWidget(self.import_config_button)

        main_layout.addLayout(export_button_layout)
        main_layout.addWidget(self.create_button)

    def create_connections(self):
        """
        Create connections between all widgets
        """

        self.set_host_ui_ctrl_button.clicked.connect(self.on_set_host_ui_button_clicked)
        self.ctrls_add_button.clicked.connect(self.on_add_ctrls_button_clicked)
        self.ctrls_remove_button.clicked.connect(self.on_remove_ctrls_button_clicked)
        self.enum_options_combo.currentIndexChanged.connect(self.on_comb_changed)
        self.export_config_button.clicked.connect(self.on_export_config_button_clicked)
        self.import_config_button.clicked.connect(self.on_import_config_button_clicked)
        self.create_button.clicked.connect(self.on_create_button_clicked)

    def on_set_host_ui_button_clicked(self):
        """
        Sets the host UI control node based on user selection.
        """

        ctrls = pm.ls(selection=True)
        if len(ctrls) != 1:
            pm.warning('Please select only one ctrl')
            return

        self._host_ui_ctrl_node = ctrls[0]
        self.host_ui_ctrl_line_edit.setText(ctrls[0].name())

    def _set_internal_data(self):
        """
        Sets the internal data based on the current state of the controls list and enum options combo.
        """

        # Set the internal data
        all_items = [self.ctrls_list.item(i).text() for i in range(self.ctrls_list.count())]
        if self.enum_options_combo.currentIndex() == 0:
            self._on_ctrls = all_items
        else:
            self._off_ctrls = all_items

    def on_comb_changed(self):
        """
        Method to handle the change event of a combo box.
        """

        # Store data and clear the list
        self.ctrls_list.clear()

        # Add all controls back to the ctrl list
        source_list = self._on_ctrls if self.enum_options_combo.currentIndex() == 0 else self._off_ctrls
        [self.ctrls_list.addItem(c) for c in source_list]

    def on_add_ctrls_button_clicked(self):
        """
        Add selected controls to a list widget and update internal data.
        """

        selected_ctrls = pm.ls(selection=True)
        for ctrl in selected_ctrls:
            if pm.nodeType(ctrl) == 'transform':
                ctrl_name = ctrl.name()

                # If we are ON and add a ctrl from the OFF list, remove it from there
                if self.enum_options_combo.currentIndex() == 0:
                    if ctrl_name in self._off_ctrls:
                        self._off_ctrls.remove(ctrl_name)
                else:
                    if ctrl_name in self._on_ctrls:
                        self._on_ctrls.remove(ctrl_name)

                # Check through the entire list to see if ctrl_name is already present
                in_list = any(self.ctrls_list.item(i).text() == ctrl_name for i in range(self.ctrls_list.count()))

                # If not present, add it to the list
                if not in_list:
                    self.ctrls_list.addItem(ctrl_name)

        # store info in the list widget
        self._set_internal_data()

    def on_remove_ctrls_button_clicked(self):
        """
        Remove the selected controls from the QListWidget.
        """

        # Get selected items in the QListWidget
        selected_list_items = [item.text() for item in self.ctrls_list.selectedItems()]

        # Iterate over the unique selected items and remove them from QListWidget
        for item_text in selected_list_items:
            item = self.ctrls_list.findItems(item_text, QtCore.Qt.MatchExactly)
            if item:
                self.ctrls_list.takeItem(self.ctrls_list.row(item[0]))

        # store info in the list widget
        self._set_internal_data()

    def generate_data(self):
        """
        Generate data for the given method.

        :return: A dictionary containing the following data:
                 - 'host_ctrl' : The name of the host control node.
                 - 'attr_name' : The text from the attribute line edit.
                 - 'on_ctrls' : A list of on controls.
                 - 'off_ctrls' : A list of off controls.
        """

        return {
            'host_ctrl': self._host_ui_ctrl_node.name(),
            'attr_name': self.attr_line_edit.text(),
            'on_ctrls': self._on_ctrls,
            'off_ctrls': self._off_ctrls
        }

    def on_export_config_button_clicked(self):
        """
        Export the configuration to a JSON file.

        Prompts the user to select a file path using a save file dialog. If a file path is selected,
        the method writes the generated dictionary data to the file as JSON.
        """

        # Open a save file dialog and get the selected path
        options = QtWidgets.QFileDialog.Options()
        fileName, _ = QtWidgets.QFileDialog.getSaveFileName(
            self,
            "QFileDialog.getSaveFileName()",
            "",
            "JSON Files (*.json);;All Files (*)",
            options=options
        )

        # If a file path is selected
        if fileName:
            # Write the dictionary to the file as json
            with open(fileName, 'w') as file:
                json.dump(self.generate_data(), file, indent=4)

    def on_import_config_button_clicked(self):
        """
        This method is called when the "Import Config" button is clicked.
        It opens a file dialog to allow the user to select a JSON file. The method then reads the
        selected file and loads it as a JSON dictionary. It validates the necessary keys exist in the
        dictionary, and if not, it displays a warning message and returns. If the keys exist,
        it sets the values of various widgets on the UI based on the values in the dictionary.
        Finally, it calls the method "on_comb_changed()".
        """

        # Open a load file dialog and get the selected path
        options = QtWidgets.QFileDialog.Options()
        fileName, _ = QtWidgets.QFileDialog.getOpenFileName(
            self,
            "QFileDialog.getOpenFileName()",
            "",
            "JSON Files (*.json);;All Files (*)",
            options=options
        )
        # if a file is selected
        if fileName:
            # read the file and load it as JSON
            with open(fileName, 'r') as file:
                config_dict = json.load(file)

            # Validate the necessary keys exist in the dictionary
            necessary_keys = ['host_ctrl', 'attr_name', 'on_ctrls', 'off_ctrls']
            if not all(key in config_dict for key in necessary_keys):
                pm.warning("Invalid config file.")
                return

            # Now set the widgets' values from the dict
            self._host_ui_ctrl_node = pm.PyNode(config_dict['host_ctrl'])
            self.host_ui_ctrl_line_edit.setText(self._host_ui_ctrl_node.name())
            self.attr_line_edit.setText(config_dict['attr_name'])
            self._on_ctrls = config_dict['on_ctrls']
            self._off_ctrls = config_dict['off_ctrls']

            self.on_comb_changed()

    def on_create_button_clicked(self):
        data = self.generate_data()
        connect_vis_attrs_from_dict(data)


if __name__ == "__main__":
    ui = VisSwitchMakerUI()
    ui.show()