ColorSetToAttribute.py

import re

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

import maya.cmds as mc
from maya.OpenMayaUI import MQtUtil

INVALID = re.compile(r'[^a-z0-9_]', re.IGNORECASE)
DIGIT = re.compile(r'[0-9]')
TYPE = {'rgb': 'vectorArray', 'rgba': 'pointArray', 'a': 'doubleArray',
        'r': 'doubleArray', 'g': 'doubleArray', 'b': 'doubleArray'}


class ColorSetToAttribute(QDialog):

    def __init__(self):
        super(ColorSetToAttribute, self).__init__(
            wrapInstance(long(MQtUtil().mainWindow()), QMainWindow))

        self.setWindowTitle('Color Set to Attribute')
        form = QFormLayout()
        self.setLayout(form)
        self.mode = QComboBox()
        form.addRow('Transfer Mode:', self.mode)
        for text in ('RGB to Vector', 'RGBA to Point', 'R to Double',
                     'G to Double', 'B to Double', 'A to Double'):
            self.mode.addItem(text)
        self.name = QLineEdit()
        form.addRow('Attribute Name:', self.name)
        self.name.textEdited.connect(self.namechanged)
        self.button = QPushButton('Transfer')
        form.addRow(self.button)
        self.button.clicked.connect(self.transfer)
        self.name.textEdited.emit('')
        self.show()

    @Slot(str)
    def namechanged(self, text):
        self.button.setEnabled(bool(text))
        self.name.setText(_validate(text))

    @Slot()
    def transfer(self):
        color_set_to_attribute(attr=self.name.text(),
                               mode=self.mode.currentText().split(' ', 1)[0])
        self.deleteLater()


def color_set_to_attribute(*meshes, **kwargs):
    attr = kwargs.get('attr')
    if not attr:
        raise RuntimeError('No attribute name specified.')
    mode = str(kwargs.get('mode', 'rgb')).lower()
    if mode not in TYPE:
        raise RuntimeError('Invalid mode specified.')
    meshes = set(meshes)
    if not meshes:
        for kwarg in ('selection', 'hilite'):
            for node in mc.ls(type='mesh', **{kwarg: True}):
                meshes.add(node)
            for node in mc.ls(type='transform', **{kwarg: True}):
                meshes.update(set(mc.listRelatives(node, type='mesh',
                                                   noIntermediate=True)))
    if not meshes:
        mc.warning('No mesh was selected.')
        return
    attr = _validate(attr)
    for mesh in meshes:
        full = mesh+'.'+attr
        try:
            lock = mc.getAttr(full, lock=True)
        except ValueError:
            lock = False
        else:
            mc.setAttr(full, lock=False)
            mc.deleteAttr(full)
        mc.addAttr(mesh, longName=attr, dataType=TYPE[mode])
        mc.setAttr(full, lock=lock)
        get = lambda x: mc.polyColorPerVertex(mesh+'.vtx[:]', q=True,
                                              **{x: True})
        if 'rgb' in mode:
            rgb = get('colorRGB')
        if 'a' in mode:
            value = get('alpha')
        elif len(mode) == 1:
            value = get('color' + mode.upper())
        if mode == 'rgb':
            value = [len(rgb)/3]
            while rgb:
                sub = []
                for i in range(3):
                    sub.append(rgb.pop(0))
                value.append(sub)
        elif mode == 'rgba':
            alpha = value
            value = [len(alpha)]
            while alpha:
                sub = []
                for j in range(3):
                    sub.append(rgb.pop(0))
                sub.append(alpha.pop(0))
                value.append(sub)
        else:
            value = [value]
        mc.setAttr(full, *value, type=TYPE[mode])

def _validate(name):
    name = INVALID.sub('_', name)
    if DIGIT.match(name):
        name = '_' + name
    return name

if __name__ == '__main__':
    ColorSetToAttribute()