Plugin to set keys on multiple attributes over multiple frames

quickKeys.py

from __future__ import absolute_import, print_function
import sys

import maya.api.OpenMaya as OpenMaya
import maya.api.OpenMayaAnim as OpenMayaAnim
from six.moves import range

maya_useNewAPI = True


class PyQuickKeysCommand(OpenMaya.MPxCommand):
    """
    Build keys (and animCurves) on multiple attributes with the provided values

    Arguments: Selection List
        A list of attributes to set data on
    -frames (-f) Time Multiple Optional
        The frames to set the keys on. If not provided, set only for the current frame
    -values (-v) Float Multiple
        The values to set to the attributes. Because of how Maya's mel commands are
        built, this must be a single flat list of values.
        If N attributes are provided, then the first N values of this list are applied
        to the attributes on the first given frame. The next N values are applied on the
        second given frame, etc...
    -uniform (-u) Bool Optional Default:False
        If True, then only take a single value per attribute and set it repeatedly
        on all the given frames
    """

    kPluginVersion = "1.0.0"
    kPluginCmdName = "quickKeys"

    kFramesFlag, kFramesFlagLong = "-f", "-frames"
    kUniformFlag, kUniformFlagLong = "-u", "-uniform"
    kValuesFlag, kValuesFlagLong = "-v", "-values"

    dependencyFn = OpenMaya.MFnDependencyNode()
    animCurveFn = OpenMayaAnim.MFnAnimCurve()

    normalCtx = OpenMaya.MDGContext.kNormal

    def __init__(self):
        super(PyQuickKeysCommand, self).__init__()

        self.times = OpenMaya.MTimeArray()
        self.values = None
        self.attrs = None
        self.curveTypes = None
        self.uniform = False

    def doIt(self, args):
        """
        Call once the first time command is run
        """
        parser = OpenMaya.MArgDatabase(self.syntax(), args)
        self.selectionList = parser.getObjectList()

        self.parseFlagArguments(parser)
        self.redoIt()

    def redoIt(self):
        """
        Call every subsequent time the command is run
        """
        self.animCurveChange = OpenMayaAnim.MAnimCurveChange()
        self.dgModifier = OpenMaya.MDGModifier()

        for i in range(self.selectionList.length()):
            plug = self.selectionList.getPlug(i)
            if not plug.isKeyable:
                continue
            if plug.isLocked:
                raise RuntimeError(
                    "Plug {0} is Locked and cannot be modified".format(plug)
                )

        for i in range(self.selectionList.length()):
            plug = self.selectionList.getPlug(i)
            if not plug.isKeyable:
                continue

            # Get animCurve connected to plug
            if OpenMayaAnim.MAnimUtil.isAnimated(plug):
                animCurve = OpenMayaAnim.MAnimUtil.findAnimation(plug)[0]
            else:
                animCurve = self.animCurveFn.create(
                    plug,
                    animCurveType=self.curveTypes[i],
                    modifier=self.dgModifier,
                )
            # Build keys on animCurve
            self.addKeys(animCurve, self.values[i])

        self.normalCtx.makeCurrent()

    def undoIt(self):
        """
        Undoes the command operations
        """
        self.animCurveChange.undoIt()
        self.dgModifier.undoIt()

    def isUndoable(self):
        """
        Indicate the command is undoable
        """
        return True

    def parseFlagArguments(self, parser):
        """
        Parse flag arguments, and format the values for setting curves
        """
        if parser.isFlagSet(self.kFramesFlag):
            self.times = self.getTimeArgs(parser, self.kFramesFlag)
        else:
            self.times.append(OpenMayaAnim.MAnimControl.currentTime())

        if parser.isFlagSet(self.kValuesFlag):
            self.values = self.getValueArgs(
                parser=parser, flagArgument=self.kValuesFlag
            )

        if parser.isFlagSet(self.kUniformFlag):
            self.uniform = parser.flagArgumentBool(self.kUniformFlag, 0)

        self.checkArguments()

        # convert the types of the value arguments
        # and group them into per-attribute chunks
        newVals = []
        self.curveTypes = []
        ta = OpenMaya.MFnUnitAttribute()
        for i in range(self.selectionList.length()):
            plug = self.selectionList.getPlug(i)
            ta.setObject(plug.attribute())
            try:
                unitType = ta.unitType()
            except RuntimeError:
                # it's an untyped attribute
                unitType = None

            # Get every Nth value from self.values for this attribute
            vals = self.values[i :: self.selectionList.length()]

            # Convert from uiUnits to internal units for distance and angle
            if unitType == OpenMaya.MFnUnitAttribute.kDistance:
                if OpenMaya.MDistance.uiUnit() != OpenMaya.MDistance.internalUnit():
                    vals = [OpenMaya.MDistance.uiToInternal(v) for v in vals]
                self.curveTypes.append(OpenMayaAnim.MFnAnimCurve.kAnimCurveTL)

            elif unitType == OpenMaya.MFnUnitAttribute.kAngle:
                if OpenMaya.MAngle.uiUnit() != OpenMaya.MAngle.internalUnit():
                    vals = [OpenMaya.MAngle.uiToInternal(v) for v in vals]
                self.curveTypes.append(OpenMayaAnim.MFnAnimCurve.kAnimCurveTA)

            else:
                self.curveTypes.append(OpenMayaAnim.MFnAnimCurve.kAnimCurveTU)

            newVals.append(vals)
        self.values = newVals

    def getTimeArgs(self, parser, flagArgument):
        """Read the -frames flag argument"""
        times = OpenMaya.MTimeArray()
        for occur in range(parser.numberOfFlagUses(self.kFramesFlag)):
            frame = parser.getFlagArgumentList(self.kFramesFlag, occur).asDouble(0)
            time = OpenMaya.MTime(frame, OpenMaya.MTime.uiUnit())
            times.append(time)
        return times

    def getValueArgs(self, parser, flagArgument):
        """Read the -value flag argument"""
        values = []
        for occur in range(parser.numberOfFlagUses(flagArgument)):
            arglist = parser.getFlagArgumentList(flagArgument, occur)
            values.append(arglist.asDouble(0))
        return values

    def checkArguments(self):
        """
        Ensure frame and transform arguments are correct in length.
        """
        if not self.values:
            raise ValueError("Values must be provided")

        if self.uniform:
            if self.selectionList.length() != len(self.values):
                raise ValueError(
                    "More than one value per attribute provided "
                    "Unable to set given uniformly throughout the frames."
                )
        else:
            requiredValues = self.selectionList.length() * len(self.times)
            if len(self.values) != requiredValues:
                msg = "{0} attributes and {1} times provided. "
                msg += "Therefore {0} * {1} = {2} values are required. Got {3}"
                raise ValueError(
                    msg.format(
                        self.selectionList.length(),
                        len(self.times),
                        requiredValues,
                        len(self.values),
                    )
                )

    def addKeys(self, animCurve, values):
        """
        Add values to given animaCurves.
        """
        self.animCurveFn.setObject(animCurve)

        # Ensures time and value length are in-sync
        if len(self.times) == 1 or self.uniform:
            uniform = values[0]
            values = OpenMaya.MDoubleArray([uniform for num in range(len(self.times))])

        # Set keys and values
        for time, value in zip(self.times, values):
            index = self.animCurveFn.find(time)
            if not index:
                # Sometimes addding key directly with the value would not trigger an
                # actually graph rebuild in Maya. This might cause the value of the key
                # to be correct, but it wont't be reflected in the channel box and the
                # user won't see the change until they manually set keys to the object
                # to trigger a graph rebuild.
                self.animCurveFn.addKey(time, value, change=self.animCurveChange)
                index = self.animCurveFn.find(time)

            self.animCurveFn.setValue(index, value, change=self.animCurveChange)

    @classmethod
    def syntaxCreator(cls):
        """
        Defines the flags and arguments of the command
        """
        syntax = OpenMaya.MSyntax()

        syntax.addFlag(cls.kFramesFlag, cls.kFramesFlagLong, syntax.kLong)
        syntax.addFlag(cls.kValuesFlag, cls.kValuesFlagLong, syntax.kLong)
        syntax.addFlag(cls.kUniformFlag, cls.kUniformFlagLong, syntax.kBoolean)

        syntax.makeFlagMultiUse(cls.kFramesFlag)
        syntax.makeFlagMultiUse(cls.kValuesFlag)

        syntax.setObjectType(syntax.kSelectionList, 1)

        return syntax

    @classmethod
    def cmdCreator(cls):
        """
        Create an instance of the command
        """
        return cls()


def initializePlugin(plugin):
    """
    Initialize the script plug-in
    """
    pluginFn = OpenMaya.MFnPlugin(
        plugin,
        "Blur",
        PyQuickKeysCommand.kPluginVersion,
        "Any",
    )

    try:
        pluginFn.registerCommand(
            PyQuickKeysCommand.kPluginCmdName,
            PyQuickKeysCommand.cmdCreator,
            PyQuickKeysCommand.syntaxCreator,
        )
    except RuntimeError:
        sys.stderr.write(
            "Failed to register command: {}\n".format(PyQuickKeysCommand.kPluginCmdName)
        )
        raise


def uninitializePlugin(plugin):
    """
    Uninitialize the script plug-in
    """
    pluginFn = OpenMaya.MFnPlugin(plugin)
    try:
        pluginFn.deregisterCommand(PyQuickKeysCommand.kPluginCmdName)
    except RuntimeError:
        sys.stderr.write(
            "Failed to unregister command: {}\n".format(
                PyQuickKeysCommand.kPluginCmdName
            )
        )
        raise