Extract Maya nParticles as Spheres

extract_particles.py

'''
Extract particles as polySpheres.

Supports:
    position, rotationPP, rgbPP, and radiusPP
'''
from maya import cmds


class Sphere(object):
    '''A class representing a colored polySphere.

    :prop position: Set/Get translation of sphere
    :prop radius: Set/Get radius of sphere
    :prop color: Set/Get color of sphere
    :meth key_position: Keyframe translation
    :meth key_radius: Keyframe radius
    :meth key_color: Keyframe color
    '''

    def __init__(self, xform, polySphere):
        self.xform = xform
        self.polySphere = polySphere

    @property
    def rotation(self):
        return cmds.xform(self.xform, q=True, ws=True, rotation=True)

    @rotation.setter
    def rotation(self, rotation):
        return cmds.xform(self.xform, ws=True, rotation=rotation)

    def key_rotation(self):
        cmds.setKeyframe(self.xform, attribute='rotate')

    def set_keys_rotation(self, keys):
        x_keys = []
        y_keys = []
        z_keys = []
        for time, value in keys:
            x_keys.append((time, value[0]))
            y_keys.append((time, value[1]))
            z_keys.append((time, value[2]))

        set_keyframes(self.xform + '.rotateX', x_keys)
        set_keyframes(self.xform + '.rotateY', y_keys)
        set_keyframes(self.xform + '.rotateZ', z_keys)

    @property
    def position(self):
        return cmds.xform(self.xform, q=True, ws=True, translation=True)

    @position.setter
    def position(self, position):
        cmds.xform(self.xform, ws=True, translation=position)

    def key_position(self):
        cmds.setKeyframe(self.xform, attribute='translate')

    def set_keys_position(self, keys):
        x_keys = []
        y_keys = []
        z_keys = []
        for time, value in keys:
            x_keys.append((time, value[0]))
            y_keys.append((time, value[1]))
            z_keys.append((time, value[2]))

        set_keyframes(self.xform + '.translateX', x_keys)
        set_keyframes(self.xform + '.translateY', y_keys)
        set_keyframes(self.xform + '.translateZ', z_keys)

    @property
    def radius(self):
        return cmds.getAttr(self.polySphere + '.radius')

    @radius.setter
    def radius(self, value):
        cmds.setAttr(self.polySphere + '.radius', value)

    def key_radius(self):
        cmds.setKeyframe(self.polySphere, attribute='radius')

    def set_keys_radius(self, keys):
        set_keyframes(self.polySphere + '.radius', keys)

    @property
    def color(self):
        return cmds.polyColorPerVertex(self.xform, query=True, rgb=True)

    @color.setter
    def color(self, color):
        cmds.polyColorSet(self.xform)
        cmds.polyColorPerVertex(self.xform, rgb=color)

    def key_color(self):
        pass

    @classmethod
    def create(cls, name, radius=0.2, color=(1, 1, 1)):
        xform, polySphere = cmds.polySphere(
            name=name,
            radius=radius
        )
        cmds.polyColorSet(
            xform,
            colorSet='rgbPP',
            create=True,
            representation='rgb'
        )
        cmds.polyColorSet(
            xform,
            colorSet='rgbPP',
            currentColorSet=True
        )
        cmds.polyColorPerVertex(xform, rgb=color)

        return cls(xform, polySphere)


def set_keyframes(attr, keys, **kwargs):
    '''Set multiple keyframes at once for an attribute.'''

    # Get MPlug for attr
    sel = om.MSelectionList()
    sel.add(attr)
    plug = sel.getPlug(0)

    # Create or get AnimCurve
    animfn = oma.MFnAnimCurve(plug)
    try:
        animfn.name()
    except:
        animfn.create(plug)

    # Set all keys
    times, values = list(zip(*keys))
    animfn.addKeys(times, values, **kwargs)


def particle_exists(shape, id):
    '''Check if a particle at id exists.'''

    return id < cmds.nParticle(shape, query=True, count=True)


def query_particle(shape, id, attr):
    '''Query a particle attribute.'''

    try:
        value = cmds.nParticle(shape, query=True, order=id, attribute=attr)
    except RuntimeError:
        value = cmds.getAttr(shape + '.' + attr[:-2])
    except:
        cmds.warning('Attribute does not exist: {}.{}'.format(shape, attr))
        raise
    return value


def get_particle_selection():
    '''Return a list of particle names and ids from your actively selected
    particles'''

    data = []
    for particle in cmds.ls(sl=True, fl=True):
        shape, id_part = particle.split('.')
        id = int(id_part.split('[')[-1].strip('[]'))
        data.append((shape, id))
    return data


def extract_particles(frames=None, color=False, rotation=False):
    '''Extracts selected particles as polySpheres'''

    cmds.undoInfo(openChunk=True)

    selection = get_particle_selection()
    spheres = [
        Sphere.create(name='{}_{}'.format(shape, id))
        for shape, id in selection
    ]

    if frames is None:
        for (sphere, (shape, id)) in zip(spheres, selection):
            sphere.position = query_particle(shape, id, 'worldPosition')
            sphere.radius = query_particle(shape, id, 'radiusPP')
            if color:
                sphere.color = query_particle(shape, id, 'rgbPP')
            if rotation:
                sphere.rotation = query_particle(shape, id, 'rotationPP')
    else:
        keys_map = {}
        for frame in frames:
            cmds.currentTime(frame, edit=True)
            for (sphere, (shape, id)) in zip(spheres, selection):

                if not particle_exists(shape, id):
                    continue

                if sphere not in keys_map:
                    keys_map[sphere] = {
                        'position': [],
                        'radius': [],
                        'color': [],
                        'rotation': [],
                    }

                keys = keys_map[sphere]
                position = query_particle(shape, id, 'worldPosition')
                keys['position'].append((frame, position))
                radius = query_particle(shape, id, 'radiusPP')
                keys['radius'].append((frame, radius))
                if color:
                    color = query_particle(shape, id, 'rgbPP')
                    keys['color'].append((frame, color))
                if rotation:
                    rotation = query_particle(shape, id, 'rotationPP')
                    keys['rotation'].append((frame, rotation))

        for sphere, keys in keys_map.items():
            sphere.set_keys_position(keys['position'])
            sphere.set_keys_radius(keys['radius'])
            if rotation:
                sphere.set_keys_rotation(keys['rotation'])

    cmds.undoInfo(closeChunk=True)
    return spheres


if __name__ == "__main__":
    extract_particles(
        frames=range(1, 120),
        color=False,
        rotation=True,
    )