Updated

Scripts/Animation/epic_pose_wrangler/v2/model/__init__.py

@@ -0,0 +1,5 @@
+from .actions import (
+    selection,
+    io,
+    zero_pose
+)

Scripts/Animation/epic_pose_wrangler/v2/model/actions/io.py

@@ -0,0 +1,68 @@
+#  Copyright Epic Games, Inc. All Rights Reserved.
+from epic_pose_wrangler.v2.model import base_action
+
+class ExportSelectedAction(base_action.BaseAction):
+    __display_name__ = "Export Selected Solvers"
+    __tooltip__ = "Exports the currently selected solver nodes in the scene to a JSON file"
+    __category__ = "IO"
+
+    @classmethod
+    def validate(cls, ui_context):
+        return bool(ui_context.current_solvers)
+
+    def execute(self, ui_context=None, **kwargs):
+        from Qt import QtWidgets
+
+        if not ui_context:
+            ui_context = self.api.get_ui_context()
+        if not ui_context:
+            return
+        file_path = QtWidgets.QFileDialog.getSaveFileName(None, "Pose Wrangler File", "", "*.json")[0]
+        # If no path is specified, exit early
+        if file_path == "":
+            return
+        self.api.serialize_to_file(file_path, ui_context.current_solvers)
+
+class ExportAllAction(base_action.BaseAction):
+    __display_name__ = "Export All Solvers"
+    __tooltip__ = "Exports all solver nodes in the scene to a JSON file"
+    __category__ = "IO"
+
+    @classmethod
+    def validate(cls, ui_context):
+        return bool(ui_context.current_solvers)
+
+    def execute(self, ui_context=None, **kwargs):
+        from Qt import QtWidgets
+
+        if not ui_context:
+            ui_context = self.api.get_ui_context()
+        if not ui_context:
+            return
+        file_path = QtWidgets.QFileDialog.getSaveFileName(None, "Pose Wrangler File", "", "*.json")[0]
+        # If no path is specified, exit early
+        if file_path == "":
+            return
+        self.api.serialize_to_file(file_path, None)
+
+class ImportFromFileAction(base_action.BaseAction):
+    __display_name__ = "Import Solvers"
+    __tooltip__ = "Imports solver nodes into the scene from a JSON file"
+    __category__ = "IO"
+
+    @classmethod
+    def validate(cls, ui_context):
+        return bool(ui_context.current_solvers)
+
+    def execute(self, ui_context=None, **kwargs):
+        from Qt import QtWidgets
+
+        if not ui_context:
+            ui_context = self.api.get_ui_context()
+        if not ui_context:
+            return
+        file_path = QtWidgets.QFileDialog.getOpenFileName(None, "Pose Wrangler File", "", "*.json")[0]
+        # If no path is specified, exit early
+        if file_path == "":
+            return
+        self.api.deserialize_from_file(file_path)

Scripts/Animation/epic_pose_wrangler/v2/model/actions/selection.py

@@ -0,0 +1,52 @@
+#  Copyright Epic Games, Inc. All Rights Reserved.
+from maya import cmds
+
+from epic_pose_wrangler.v2.model import base_action
+
+class SelectSolverAction(base_action.BaseAction):
+    __display_name__ = "Select Solver Node(s)"
+    __tooltip__ = "Selects the currently selected solver nodes in the scene"
+    __category__ = "Select"
+
+    @classmethod
+    def validate(cls, ui_context):
+        return bool(ui_context.current_solvers)
+
+    def execute(self, ui_context=None, **kwargs):
+        if not ui_context:
+            ui_context = self.api.get_ui_context()
+        if not ui_context:
+            return
+        cmds.select(ui_context.current_solvers, replace=True)
+
+class SelectDriverAction(base_action.BaseAction):
+    __display_name__ = "Select Driver Node(s)"
+    __tooltip__ = "Selects the driver nodes in the scene"
+    __category__ = "Select"
+
+    @classmethod
+    def validate(cls, ui_context):
+        return bool(ui_context.drivers)
+
+    def execute(self, ui_context=None, **kwargs):
+        if not ui_context:
+            ui_context = self.api.get_ui_context()
+        if not ui_context:
+            return
+        cmds.select(ui_context.drivers, replace=True)
+
+class SelectDrivenAction(base_action.BaseAction):
+    __display_name__ = "Select Driven Node(s)"
+    __tooltip__ = "Selects the driven nodes in the scene"
+    __category__ = "Select"
+
+    @classmethod
+    def validate(cls, ui_context):
+        return bool(ui_context.driven)
+
+    def execute(self, ui_context=None, **kwargs):
+        if not ui_context:
+            ui_context = self.api.get_ui_context()
+        if not ui_context:
+            return
+        cmds.select(ui_context.driven, replace=True)

Scripts/Animation/epic_pose_wrangler/v2/model/actions/zero_pose.py

@@ -0,0 +1,41 @@
+#  Copyright Epic Games, Inc. All Rights Reserved.
+from maya import cmds
+from epic_pose_wrangler.v2.model import base_action, pose_blender
+
+class ZeroDefaultPoseAction(base_action.BaseAction):
+    __display_name__ = "Zero Default Pose Transforms"
+    __tooltip__ = ""
+    __category__ = "Utilities"
+
+    @classmethod
+    def validate(cls, ui_context):
+        return bool(ui_context.current_solvers)
+
+    def execute(self, ui_context=None, solver=None, **kwargs):
+        from Qt import QtWidgets
+        if not ui_context:
+            ui_context = self.api.get_ui_context()
+        if ui_context:
+            solver = self.api.get_rbf_solver_by_name(ui_context.current_solvers[-1])
+        if not solver:
+            solver = self.api.current_solver
+
+        # Go to the default pose
+        solver.go_to_pose('default')
+        # Assume edit is enabled
+        edit = True
+        if not self.api.get_solver_edit_status(solver):
+            # If edit isn't enabled, store the current enabled state and enable editing
+            edit = False
+            self.api.edit_solver(edit=True, solver=solver)
+
+        # Reset the driven transforms
+        for node in solver.driven_nodes(type=pose_blender.UEPoseBlenderNode.node_type):
+            cmds.setAttr('{node}.translate'.format(node=node), 0.0, 0.0, 0.0)
+            cmds.setAttr('{node}.rotate'.format(node=node), 0.0, 0.0, 0.0)
+            cmds.setAttr('{node}.scale'.format(node=node), 1.0, 1.0, 1.0)
+
+        # Update the pose
+        self.api.update_pose(pose_name='default', solver=solver)
+        # Restore edit status
+        self.api.edit_solver(edit=edit, solver=solver)

Scripts/Animation/epic_pose_wrangler/v2/model/base_action.py

@@ -0,0 +1,23 @@
+#  Copyright Epic Games, Inc. All Rights Reserved.
+import abc
+
+class BaseAction(object):
+    __display_name__ = "BaseAction"
+    __tooltip__ = ""
+    __category__ = ""
+
+    @classmethod
+    @abc.abstractmethod
+    def validate(cls, ui_context):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def execute(self, ui_context=None, **kwargs):
+        raise NotImplementedError
+
+    def __init__(self, api=None):
+        self._api = api
+
+    @property
+    def api(self):
+        return self._api

Scripts/Animation/epic_pose_wrangler/v2/model/base_extension.py

@@ -0,0 +1,58 @@
+#  Copyright Epic Games, Inc. All Rights Reserved.
+from epic_pose_wrangler.model import exceptions
+
+class PoseWranglerExtension(object):
+    """
+    Base class for extending pose wrangler with custom utilities that can be dynamically added to the UI
+    """
+    __category__ = ""
+
+    def __init__(self, display_view=False, api=None):
+        super(PoseWranglerExtension, self).__init__()
+        self._display_view = display_view
+        self._view = None
+        self._api = api
+
+    @property
+    def api(self):
+        """
+        Get the current API
+
+        :return: Reference to the main API interface
+        :rtype: pose_wrangler.v2.main.UERBFAPI
+        """
+
+        return self._api
+
+    @property
+    def view(self):
+        """
+        Get the current view widget. This should be overridden by custom extensions if you wish to embed a UI for this
+        extension into the main PoseWrangler UI.
+
+        :return: Reference to the PySide widget associated with this extension
+        :rtype: QWidget or None
+        """
+        return self._view
+
+    def execute(self, context=None, **kwargs):
+        """
+        Generic entrypoint for executing the extension.
+
+        :param: context: pose wrangler context containing current solver and all solvers
+        :type context: pose_wrangler.v2.model.context.PoseWranglerContext or None
+        """
+        raise exceptions.PoseWranglerFunctionalityNotImplemented(
+            "'execute' function has not been implemented for {class_name}".format(
+                class_name=self.__class__.__name__
+            )
+        )
+
+    def on_context_changed(self, new_context):
+        """
+        Context event called when the current solver is set via the API
+
+        :param new_context: pose wrangler context containing current solver and all solvers
+        :type new_context: pose_wrangler.v2.model.context.PoseWranglerContext or None
+        """
+        pass

Scripts/Animation/epic_pose_wrangler/v2/model/context.py

@@ -0,0 +1,13 @@
+#  Copyright Epic Games, Inc. All Rights Reserved.
+class PoseWranglerContext(object):
+    def __init__(self, current_solver, solvers):
+        self._current_solver = current_solver
+        self._solvers = solvers
+
+    @property
+    def current_solver(self):
+        return self._current_solver
+
+    @property
+    def solvers(self):
+        return self._solvers

Scripts/Animation/epic_pose_wrangler/v2/model/exceptions.py

@@ -0,0 +1,45 @@
+#  Copyright Epic Games, Inc. All Rights Reserved.
+"""
+API v2 Specific exceptions
+"""
+from epic_pose_wrangler.model import exceptions
+
+class MessageConnectionError(exceptions.PoseWranglerException):
+    """
+    Raised when message connections fail.
+    """
+
+class InvalidSolverError(exceptions.PoseWranglerException):
+    """
+    Raised when the incorrect solver type is specified
+    """
+
+class InvalidNodeType(exceptions.PoseWranglerException, TypeError):
+    """
+    Raised when the incorrect node type is specified
+    """
+
+class PoseWranglerAttributeError(exceptions.PoseWranglerException, AttributeError):
+    """
+    Raised when there is an issue getting/setting an attribute
+    """
+
+class PoseBlenderPoseError(exceptions.PoseWranglerException):
+    """
+    Generic error for issues with poses
+    """
+
+class InvalidPose(exceptions.PoseWranglerException):
+    """
+    Generic error for incorrect poses
+    """
+
+class InvalidPoseIndex(exceptions.PoseWranglerException):
+    """
+    Raised when issues arise surrounding the poses index
+    """
+
+class BlendshapeError(exceptions.PoseWranglerException):
+    """
+    Generic error for blendshape issues
+    """

Scripts/Animation/epic_pose_wrangler/v2/model/export.py

@@ -0,0 +1,325 @@
+#  Copyright Epic Games, Inc. All Rights Reserved.
+import os
+import json
+from maya import cmds
+
+from special_projects.publish_tools.fbx_cmd import fbx_export
+from special_projects.publish_tools.utils import find_top_joints
+
+from special_projects.rigging.rbf_node import RBFNode
+
+class RBFNodeExporter(object):
+    """
+    Utility class to export a RBFsolver node, exports a JSON and FBX
+    
+    >>> from special_projects.rigging.rbf_node.export import RBFNodeExporter
+    >>> node = 'my_UE4RBFSolver_node'
+    >>> asset_name = 'my_asset_name'
+    >>> export_directory = 'my:/export/directory'
+    >>> exporter = RBFNodeExporter(node, asset_name, export_directory)
+    >>> exporter.export()
+
+    Result:
+        export_directory/node.json
+        export_directory/node.fbx
+    """
+
+    def __init__(self, node, asset_name, export_directory):
+        """
+        Initialize RBF Exporter
+        """
+
+        self.set_node(node)
+        self.set_asset_name(asset_name)
+        self.set_export_directory(export_directory)
+        self._set_fbx_export_path()
+        self._set_json_export_path()
+
+    def node(self):
+        return self._node
+
+    def set_node(self, node):
+
+        if not cmds.objectType(node, isAType=RBFNode.node_type):
+            raise TypeError('Invalid "{}" node: "{}"'.format(RBFNode.node_type, node))
+
+        self._node = RBFNode(node)
+
+    def asset_name(self):
+        return self._asset_name
+
+    def set_asset_name(self, asset_name):
+        self._asset_name = asset_name
+
+    def export_directory(self):
+        return self._export_directory
+
+    def set_export_directory(self, directory):
+        if os.path.isdir(directory) and os.path.exists(directory):
+            self._export_directory = directory
+        else:
+            raise IOError('Export directory "{}" does not exists'.format(directory))
+
+    def fbx_export_path(self):
+        return self._fbx_export_path
+
+    def _set_fbx_export_path(self):
+        self._fbx_export_path = os.path.join(self._export_directory, '{}.fbx'.format(self._node))
+
+    def json_export_path(self):
+        return self._json_export_path
+
+    def _set_json_export_path(self):
+        self._json_export_path = os.path.join(self._export_directory, '{}.json'.format(self._node))
+
+    # -------------------------------------------------------------------------------------
+
+    def export(self):
+        """
+        Exports FBX and sidecar JSON file
+        """
+
+        self.fbx_export()
+        self.json_export()
+
+    def json_export(self):
+        """
+        Exports JSON sidecar
+        """
+
+        self.node().name_unnamed_poses()
+
+        config = self._node.data()
+        config['export_fbx'] = self.fbx_export_path()
+        config['asset_name'] = self.asset_name()
+
+        with open(self.json_export_path(), 'w') as outfile:
+            json.dump(config, outfile, sort_keys=0, indent=4, separators=(",", ":"))
+
+    def fbx_export(self):
+        """
+        Exports baked poses to FBX
+        """
+
+        self.node().name_unnamed_poses()
+
+        self.bake_poses()
+        cmds.select(self.root_joint())
+        fbx_export(self.fbx_export_path(),
+                   animation=True,
+                   bake_complex_animation=True,
+                   bake_complex_start=0,
+                   bake_complex_end=cmds.playbackOptions(q=True, maxTime=True),
+                   up_axis='z')
+
+    # -------------------------------------------------------------------------------------
+
+    def blendshape_nodes(self):
+        """
+        Finds blendshape nodes from driven attributes
+        """
+
+        driven_attributes = self._node.driven_attributes(type='blendShape')
+
+        blendshape_nodes = list()
+        for output in driven_attributes:
+            for attribute in output:
+                blendshape_nodes.append(attribute.split('.')[0])
+
+        return list(set(blendshape_nodes))
+
+    def meshes(self):
+        """
+        Finds meshes from blendshape nodes
+        """
+
+        meshes = list()
+
+        blendShapes = self.blendshape_nodes()
+        if blendShapes:
+            for blendShape in blendShapes:
+                meshes.extend(cmds.deformer(blendShape, q=True, geometry=True))
+
+            meshes = list(set(meshes))
+            meshes = cmds.listRelatives(meshes, parent=True)
+
+        return meshes
+
+    def root_joint(self):
+        """
+        Finds root joint from meshes
+        """
+
+        meshes = self.meshes()
+        if meshes:
+            skeleton = list()
+            for mesh in meshes:
+                skin_cluster = cmds.ls(cmds.findDeformers(mesh), type='skinCluster')
+                if skin_cluster:
+                    skin_cluster = skin_cluster[0]
+                    influences = cmds.skinCluster(skin_cluster, q=True, inf=True)
+                    if influences:
+                        skeleton.extend(influences)
+                    else:
+                        raise RuntimeError('No influences found for "{}"'.format(skin_cluster))
+                else:
+                    cmds.warning('No skinCluster found for "{}"'.format(mesh))
+
+            root_joints = find_top_joints(skeleton)
+
+        else:
+            skeleton = self._node.drivers()
+            # for driver in self._node.drivers():
+            # driven = driver.replace('_drv', '')
+            # if cmds.objExists(driven):
+            #    skeleton.append(driven)
+
+            root_joints = find_top_joints(skeleton)
+
+        if not root_joints:
+            raise RuntimeError('No root joint found for "{}"'.format(self._node))
+
+        return root_joints[0]
+
+    def add_root_attributes(self, root_joint):
+        """
+        Adds RBFNode driven attributes to root_joint
+        """
+
+        pose_root_attributes = dict()
+
+        poses = self._node.poses()
+        for pose in poses:
+            driven_attributes = self._node.pose_driven_attributes(pose)  # add type flag!
+            current_pose = list()
+            for attribute in driven_attributes:
+                node, target = attribute.split('.')
+                root_attribute = '{}.{}'.format(root_joint, target)
+                if root_attribute not in current_pose:
+                    if not cmds.objExists(root_attribute):
+                        cmds.addAttr(root_joint, ln=target, at='double', k=True)
+                    else:
+                        input_connection = cmds.listConnections(root_attribute, s=True, d=False, plugs=True)
+                        if input_connection:
+                            cmds.disconnectAttr(input_connection[0], root_attribute)
+
+                    # cmds.connectAttr(attribute, root_attribute)
+                    current_pose.append(root_attribute)
+
+            pose_root_attributes[pose] = current_pose
+
+        return pose_root_attributes
+
+    def bake_poses(self):
+        """
+        Bakes the RBFNode poses in the timeline for FBX export
+        """
+
+        for anim_curve_type in ['animCurveTL', 'animCurveTA', 'animCurveTU']:
+            cmds.delete(cmds.ls(type=anim_curve_type))
+
+        pose_root_attributes = self.add_root_attributes(self.root_joint())
+
+        for frame, pose in enumerate(self._node.poses()):
+
+            # go to pose
+            self._node.go_to_pose(pose)
+
+            # key controllers
+            if self._node.num_controllers():
+                for controller in self._node.controllers():
+                    cmds.setKeyframe(controller, t=frame, inTangentType='linear', outTangentType='step')
+
+            # or key drivers
+            else:
+                for driver in self._node.drivers():
+                    cmds.setKeyframe(driver, t=frame, inTangentType='linear', outTangentType='step')
+
+            root_attributes = pose_root_attributes.get(pose, [])
+
+            for root_attribute in root_attributes:
+
+                input_connection = cmds.listConnections(root_attribute, s=True, d=False, plugs=True)
+                if input_connection:
+                    cmds.disconnectAttr(input_connection[0], root_attribute)
+
+                # Key Driven Before/After
+                cmds.setAttr(root_attribute, 0)
+
+                if frame == len(self._node.poses()) - 1:
+                    cmds.setKeyframe(root_attribute, t=(frame - 1), inTangentType='linear', outTangentType='linear')
+                else:
+                    cmds.setKeyframe(root_attribute, t=((frame - 1), (frame + 1)), inTangentType='linear',
+                                     outTangentType='linear')
+
+                # Key Driven
+                cmds.setAttr(root_attribute, 1)
+                cmds.setKeyframe(root_attribute, t=frame, inTangentType='linear', outTangentType='linear')
+
+        # set start-end frames
+        end_frame = len(self._node.poses()) - 1
+        cmds.playbackOptions(minTime=0, maxTime=end_frame, animationStartTime=0, animationEndTime=end_frame)
+        cmds.dgdirty(a=True)
+
+class RBFPoseExporterBatch(object):
+    """
+    Utility class to export multiple RBFsolver nodes, exports a JSON and FBX for each solver
+
+    >>> solvers = cmds.ls(type='UE4RBFSolverNode')
+    >>> rig_scene = r'D:\Build\UE5_Main\Collaboration\Frosty\ArtSource\Character\Hero\Kenny\Rig\Kenny_Rig.ma'
+    >>> export_directory = r'D:\test\class'
+    >>> asset_name = 'Kenny'
+    >>> rbf_exporter_batch = RBFPoseExporterBatch(solvers, asset_name, export_directory, rig_scene)
+
+    Result:
+        export_directory/node1.json
+        export_directory/node1.fbx
+        export_directory/node2.json
+        export_directory/node2.fbx
+    """
+
+    def __init__(self, nodes, asset_name, export_directory, rig_scene):
+        self.set_pose_exporter(nodes, asset_name, export_directory)
+        self.set_rig_scene(rig_scene)
+
+    def pose_exporter(self):
+        return self._poseExporter
+
+    def asset_name(self):
+        return self._asset_name
+
+    def export_directory(self):
+        return self._export_directory
+
+    def set_pose_exporter(self, nodes, asset_name, export_directory):
+        if not hasattr(nodes, '__iter__'):
+            nodes = [nodes]
+
+        exporters = list()
+        for node in nodes:
+            if cmds.objectType(node, isAType=RBFNodeExporter.node_type):
+                exporters.append(RBFNodeExporter(node, asset_name, export_directory))
+
+        if not len(exporters):
+            raise RuntimeError('No valid {} objects found'.format(RBFNodeExporter.node_type))
+
+        self._poseExporter = exporters
+        self._asset_name = asset_name
+        self._export_directory = export_directory
+
+    def rig_scene(self):
+        return self._rig_scene
+
+    def set_rig_scene(self, rig_scene):
+        if not os.path.exists(rig_scene):
+            raise IOError('Rig scene "{}" does not exists'.format(rig_scene))
+
+        self._rig_scene = rig_scene
+
+    # -------------------------------------------------------------------------------------
+
+    def export(self, run_in_subprocess=True):
+        # TO-DO: Implement run_in_subprocess
+        for exporter in self.pose_exporter():
+            cmds.file(self.rig_scene(), open=True, force=True)
+            exporter.export()

Scripts/Animation/epic_pose_wrangler/v2/model/pose_blender.py

@@ -0,0 +1,501 @@
+#  Copyright Epic Games, Inc. All Rights Reserved.
+
+import copy
+
+from maya import cmds
+from maya.api import OpenMaya as om
+
+from epic_pose_wrangler.log import LOG
+from epic_pose_wrangler.v2.model import exceptions, utils
+
+class UEPoseBlenderNode(object):
+    """
+    Class wrapper for UEPoseBlenderNode
+    """
+    node_type = 'UEPoseBlenderNode'
+
+    @classmethod
+    def create(cls, driven_transform=None):
+        """
+        Create Pose Blender Node
+        >>> new_node = UEPoseBlenderNode.create()
+        """
+        name = driven_transform
+        # If no name is specified, generate unique name
+        if name is None:
+            name = '{name}#'.format(name=cls.node_type)
+        name += "_{class_name}".format(class_name=cls.__name__)
+        # Create node
+        node = cmds.createNode(cls.node_type, name=name)
+        node_ref = cls(node)
+        # If a driving transform is specified, connect it via a message attribute and set the base pose
+        if driven_transform is not None:
+            utils.message_connect(
+                from_attribute="{node}.drivenTransform".format(node=node),
+                to_attribute="{transform}.poseBlender".format(transform=driven_transform)
+            )
+            # Set the base pose
+            node_ref.base_pose = cmds.xform(driven_transform, query=True, objectSpace=True, matrix=True)
+        # Return the new UEPoseBlenderNode reference
+        return node_ref
+
+    @classmethod
+    def find_all(cls):
+        """
+        Returns all PoseBlender nodes in scene
+        """
+        return [cls(node) for node in cmds.ls(type=cls.node_type)]
+
+    @classmethod
+    def find_by_name(cls, name):
+        """
+        Find a UEPoseBlenderNode class with the specified name
+        :param name :type str: name of the node
+        :return :type UEPoseBlenderNode or None:
+        """
+        if not name.endswith(cls.node_type):
+            name += "_{node_type}".format(node_type=cls.node_type)
+        match = cmds.ls(name, type=cls.node_type)
+        return cls(match[0]) if match else None
+
+    @classmethod
+    def find_by_transform(cls, transform):
+        """
+        Finds a UEPoseBlenderNode class connected to the specified joint
+        :param transform :type str: name of the transform
+        :return :type UEPoseBlenderNode or None
+        """
+        if not cmds.attributeQuery('poseBlender', node=transform, exists=True):
+            return
+        connections = cmds.listConnections('{transform}.poseBlender'.format(transform=transform))
+        if not connections:
+            return
+        return cls(connections[0])
+
+    def __init__(self, node):
+        # If the node doesn't exist, raise an exception
+        if not cmds.objectType(node, isAType=self.node_type):
+            raise exceptions.InvalidNodeType(
+                'Invalid "{node_type}" node: "{node_name}"'.format(
+                    node_type=self.node_type,
+                    node_name=node
+                )
+            )
+        # Store a reference ot the MObject in case the name of the node is changed whilst this class is still in use
+        self._node = om.MFnDependencyNode(om.MGlobal.getSelectionListByName(node).getDependNode(0))
+
+    def __repr__(self):
+        """
+        Returns class string representation
+        """
+        return '<{node_type}>: {object}'.format(node_type=self.node_type, object=self)
+
+    def __str__(self):
+        """
+        Returns class as string
+        """
+        return str(self.node)
+
+    def __eq__(self, other):
+        """
+        Returns if two objects are the same, allows for comparing two different UEPoseBlenderNode references that wrap
+        the same MObject
+        """
+        return str(self) == str(other)
+
+    def __enter__(self):
+        """
+        Override the __enter__ to allow for this class to be used as a context manager to toggle edit mode
+        >>> pose_blender = UEPoseBlenderNode('TestPoseBlenderNode')
+        >>> with pose_blender:
+        >>>     # Edit mode enabled so changes can be made
+        >>>     pass
+        >>> # Edit mode is now disabled
+        """
+        self.edit = True
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        On exit, disable edit mode
+        """
+        self.edit = not self.edit
+
+    @property
+    def node(self):
+        return self._node.name()
+
+    @property
+    def rbf_solver_attr(self):
+        return '{node}.rbfSolver'.format(node=self.node)
+
+    @property
+    def base_pose_attr(self):
+        return '{node}.basePose'.format(node=self.node)
+
+    @property
+    def envelope_attr(self):
+        return '{node}.envelope'.format(node=self.node)
+
+    @property
+    def in_matrix_attr(self):
+        return '{node}.inMatrix'.format(node=self.node)
+
+    @property
+    def out_matrix_attr(self):
+        return '{node}.outMatrix'.format(node=self.node)
+
+    @property
+    def poses_attr(self):
+        return '{node}.poses'.format(node=self.node)
+
+    @property
+    def weights_attr(self):
+        return '{node}.weights'.format(node=self.node)
+
+    @property
+    def driven_transform(self):
+        """
+        Get the current driven transform
+        :return :type str: transform node name
+        """
+        if cmds.attributeQuery('drivenTransform', node=self.node, exists=True):
+            transforms = cmds.listConnections("{node}.drivenTransform".format(node=self.node))
+            if transforms:
+                return transforms[0]
+
+    @property
+    def edit(self):
+        """
+        :return :type bool: are the driven transforms connected to this node editable?
+        """
+        transform = self.driven_transform
+        # If no transform, we can edit!
+        if not transform:
+            return True
+        # False if we have any matrix connections, otherwise we have no connections and can edit
+        return False if cmds.listConnections(self.out_matrix_attr) else True
+
+    @edit.setter
+    def edit(self, value):
+        """
+        Allow the driven transforms to be edited (or not)
+        :param value :type bool: edit mode enabled True/False
+        """
+        # If we don't have a driven transform we can't do anything
+        transform = self.driven_transform
+        if not transform:
+            return
+        # If we are in edit mode, disable out connections
+        if bool(value):
+            self.out_matrix = None
+        # Otherwise connect up the transform
+        else:
+            self.out_matrix = transform
+
+    @property
+    def rbf_solver(self):
+        """
+        :return:The connected rbf solver node if it exists
+        """
+        from epic_pose_wrangler.v2.model import api
+        # Check the attr exists
+        exists = cmds.attributeQuery(self.rbf_solver_attr.split('.')[-1], node=self.node, exists=True)
+        # Query the connections
+        connections = cmds.listConnections(self.rbf_solver_attr)
+        # Return the node if the attr exists and a connection is found
+        return api.RBFNode(connections[0]) if exists and connections else None
+
+    @rbf_solver.setter
+    def rbf_solver(self, rbf_solver_attr):
+        """
+        Connects up this node to an rbf solver node
+        :param rbf_solver_attr: the attribute on the rbf solver to connect this to i.e my_rbf_solver.poseBlend_back_120
+        """
+        utils.message_connect(rbf_solver_attr, self.rbf_solver_attr)
+
+    @property
+    def base_pose(self):
+        """
+        :return :type list: matrix
+        """
+        return self.get_pose(index=0)
+
+    @base_pose.setter
+    def base_pose(self, matrix):
+        """
+        Set the base pose to a specific pose.PoseNode
+        :param matrix :type list: Matrix to set as base pose
+        """
+        # Connect up the pose's outputLocalMatrix to the basePose plug
+        utils.set_attr_or_connect(source_attr_name=self.base_pose_attr, value=matrix, attr_type='matrix')
+        # Set the inMatrix to the basePose plug
+        self.in_matrix = matrix
+        # Set the first pose in the pose list to the new base pose
+        self.set_pose(index=0, overwrite=True, matrix=matrix)
+
+    @property
+    def envelope(self):
+        """
+        Get the value of the envelope attr from this node
+        """
+        return utils.get_attr(self.envelope_attr, as_value=True)
+
+    @envelope.setter
+    def envelope(self, value):
+        """
+        Sets the envelope
+        :param value: float, int or string (i.e node.attributeName). Float/Int will set the value, whilst
+        passing an attribute will connect the plugs
+        """
+        if isinstance(value, float) or isinstance(value, int) and 0 > value > 1:
+            value = min(max(0.0, value), 1.0)
+        utils.set_attr_or_connect(source_attr_name=self.envelope_attr, value=value)
+
+    @property
+    def in_matrix(self):
+        """
+        Get the inMatrix value
+        :return: matrix or attributeName
+        """
+        # Get the current value of the attribute
+        return utils.get_attr(self.in_matrix_attr, as_value=True)
+
+    @in_matrix.setter
+    def in_matrix(self, value):
+        """
+        Sets the inMatrix attribute to a value or connects it to another matrix plug
+        :param value: matrix (list) or string node.attributeName
+        """
+        utils.set_attr_or_connect(source_attr_name=self.in_matrix_attr, value=value, attr_type='matrix')
+
+    @property
+    def out_matrix(self):
+        """
+        :return:
+        """
+        return cmds.getAttr(self.out_matrix_attr)
+
+    @out_matrix.setter
+    def out_matrix(self, transform_name=None):
+        """
+        Connect the output matrix up to a given transform node
+        :param transform_name: name of a transform node to connect to
+        """
+        # If a transform is specified, connect it
+        current_selection = cmds.ls(selection=True)
+        if transform_name:
+            # Check that the node type is a transform
+            if not cmds.ls(transform_name, type='transform'):
+                # If its not a transform raise an error
+                raise exceptions.InvalidNodeType(
+                    "Invalid node type. Expected 'transform', received '{node_type}'".format(
+                        node_type=cmds.objectType(transform_name)
+                    )
+                )
+            # Create a new decomposeMatrix node to convert the outMatrix to T,R,S
+            mx_decompose_node = cmds.createNode('decomposeMatrix', name="{node}_mx_decompose#".format(node=self.node))
+            # Connect the outMatrix up to the decomposeMatrix's input
+            cmds.connectAttr(self.out_matrix_attr, '{node}.inputMatrix'.format(node=mx_decompose_node))
+            # Create an iterator with TRS attributes
+            attributes = ('translate', 'rotate', 'scale')
+            # Iterate through each attr and connect it
+            for attr in attributes:
+                out_attr = '{mx_decompose_node}.output{attr}'.format(
+                    mx_decompose_node=mx_decompose_node,
+                    attr=attr.capitalize()
+                )
+                in_attr = '{node_name}.{attr}'.format(node_name=transform_name, attr=attr)
+                cmds.connectAttr(out_attr, in_attr, force=True)
+        else:
+            # No transform was specified, disconnect all existing connections
+            connections = cmds.listConnections(self.out_matrix_attr, type='decomposeMatrix')
+            current_matrix = cmds.xform(self.driven_transform, query=True, matrix=True, objectSpace=True)
+            if connections:
+                cmds.delete(connections)
+            for connection in cmds.listConnections(self.out_matrix_attr, plugs=True) or []:
+                cmds.disconnectAttr(self.out_matrix_attr, connection)
+            cmds.xform(self.driven_transform, matrix=current_matrix, objectSpace=True)
+        cmds.select(current_selection, replace=True)
+
+    def get_pose(self, index=-1):
+        """
+        Get the pose at the specified index
+        :param index :type int: index to query
+        :return :type list matrix or None
+        """
+        # Get the value of the pose_attr at the specified index
+        return utils.get_attr(
+            '{poses_attr}[{index}]'.format(poses_attr=self.poses_attr, index=index),
+            as_value=True
+        )
+
+    def get_poses(self):
+        """
+        :return :type list of matrices
+        """
+        # TODO update this when the weight + name get added
+        return utils.get_attr_array(attr_name=self.poses_attr, as_value=True)
+
+    def add_pose_from_current(self, pose_name, index=-1):
+        """
+        Add a pose from the current transforms positions
+        :param pose_name :type str: name of the pose
+        :param index :type int: target index for the pose
+        """
+        # Find the next index if no index is specified
+        if index < 0:
+            # If the index is less than 0 find the next available index
+            index = len(cmds.getAttr(self.poses_attr, multiIndices=True) or [0]) - 1
+        # Store driven transform in var to reduce number of cmds calls
+        driven_transform = self.driven_transform
+        if not driven_transform:
+            raise exceptions.PoseBlenderPoseError(
+                "No driven transform associated with this node, "
+                "unable to get the current matrix"
+            )
+        # Get the local matrix for the driven transform
+        local_matrix = cmds.xform(driven_transform, query=True, objectSpace=True, matrix=True)
+        # Set the pose
+        self.set_pose(index=index, pose_name=pose_name, matrix=local_matrix)
+
+    def set_pose(self, index=-1, pose_name="", overwrite=True, matrix=None):
+        """
+        Set a pose for the specified index
+        :param index: int value of the index to set
+        :param pose_name: name of the pose as a string
+        :param overwrite: should it overwrite any existing pose at the index or insert
+        :param matrix: matrix value to set
+        """
+        # Find the next index if no index is specified
+        if index < 0:
+            if pose_name:
+                LOG.warning("Set Pose has not been implemented to support a pose name. FIXME")
+                return
+            else:
+                # If the index is less than 0 find the next available index
+                index = len(cmds.getAttr(self.poses_attr, multiIndices=True) or [0]) - 1
+
+        # If no matrix is specified, grab the matrix for the driven transform
+        if matrix is None:
+            matrix = cmds.xform(self.driven_transform, query=True, matrix=True, objectSpace=True)
+
+        # TODO add in support for inserting into the pose list
+        if not overwrite:
+            pose_count = len(cmds.getAttr(self.poses_attr, multiIndices=True) or [])
+            if pose_count - 1 > index:
+                pass
+
+        # Generate the source attribute (the attribute on this node) from the default attr and given index
+        source_attr_name = '{poses_attr}[{index}]'.format(poses_attr=self.poses_attr, index=index)
+        # Set the pose
+        utils.set_attr_or_connect(source_attr_name=source_attr_name, value=matrix, attr_type='matrix')
+
+    def go_to_pose(self, index=-1):
+        """
+        Move the driven transform to the matrix stored in the pose list at the specified index
+        :param index :type index: index to go to
+        """
+        # Get the matrix at the specified index
+        pose_matrix = self.get_pose(index=index)
+        # Assume the position
+        cmds.xform(self.driven_transform, matrix=pose_matrix)
+
+    def delete_pose(self, index=-1, pose_name=""):
+        """
+        Delete a pose at the specified index or with the given name
+        :param index :type int: index to delete or -1 to use pose_name
+        :param pose_name :type str: pose name to delete
+        """
+        # TODO Delete pose will need to remember enabled state + reconnect up pose_name connection when the changes are
+        #  made to the solver. Currently can't delete by pose name because there is nothing tying an index to a pose
+        #  name
+        # If no index and no pose name, raise exception
+        if index < 0 and not pose_name:
+            raise exceptions.InvalidPoseIndex("Unable to delete pose")
+
+        # Copy a list of the poses
+        poses = copy.deepcopy(self.get_poses())
+
+        # Iterate through the existing poses in reverse
+        for i in reversed(range(len(poses))):
+            # Remove the pose from the list of targets
+            cmds.removeMultiInstance('{blender}.poses[{index}]'.format(blender=self, index=i), b=True)
+
+        # Remove the pose at the given index
+        poses.pop(index)
+
+        # Re-add all the deleted poses (minus the one we popped)
+        for pose_index, matrix in enumerate(poses):
+            self.set_pose(index=pose_index, matrix=matrix)
+
+    def get_weight(self, index, as_float=True):
+        """
+        Get the current weight for the specified index
+        :param index :type int
+        :param as_float :type bool: return as float or as attribute name of the connected plug
+        :return: float or plug
+        """
+        return utils.get_attr('{weights}[{index}]'.format(weights=self.weights_attr, index=index), as_value=as_float)
+
+    def get_weights(self, as_float=True):
+        """
+        Get all the weights associated with this node
+        :param as_float :type bool: as list of floats or list of plugs
+        :return: list of floats or plugs
+        """
+        return utils.get_attr_array(attr_name=self.weights_attr, as_value=as_float)
+
+    def set_weight(self, index=-1, in_float_attr="", float_value=0.0):
+        """
+        Set the weight at a given index either by connecting an attribute or specifying a float value
+        :param index :type int: index to set
+        :param in_float_attr :type str: node.attributeName to connect to this plug
+        :param float_value :type float: float value to set
+        """
+        if index < 0:
+            # If the index is less than 0 find the next available index
+            index = len(cmds.getAttr(self.weights_attr, multiIndices=True) or [0]) - 1
+        # Generate the correct source attr name for the index specified
+        source_attr_name = '{weights}[{index}]'.format(weights=self.weights_attr, index=index)
+        # Set the array element to either the attr or the float value
+        utils.set_attr_or_connect(source_attr_name=source_attr_name, value=in_float_attr or float_value)
+
+    def set_weights(self, in_float_array_attr="", floats=None):
+        """
+        Set multiple weights either by an attribute array string or a float list.Will overwrite existing values
+        :param in_float_array_attr :type str: node.attributeName array attribute to connect all indices with
+        :param floats :type list: list of float values to set for each corresponding index
+        """
+        # Prioritize plugs over setting floats
+        if in_float_array_attr:
+            # Iterate through all the indices in the array
+            for i in range(len(cmds.getAttr(in_float_array_attr, multiIndices=True) or [0])):
+                # Generate the source attr name
+                in_float_attr = "{array_attr}[{index}]".format(array_attr=in_float_array_attr, index=i)
+                # Set the weight for the current index
+                self.set_weight(index=i, in_float_attr=in_float_attr)
+        elif floats:
+            # Iterate through the floats
+            for index, float_value in enumerate(floats):
+                # Set the weight for the current index
+                self.set_weight(index=index, float_value=float_value)
+
+    def delete(self):
+        """
+        Delete the node associated with this wrapper
+        """
+        # If we aren't in edit mode we still have connections made to decomposeMatrix nodes that we want to delete
+        if not self.edit:
+            # Enable edit mode to delete those decomposeMatrix nodes
+            self.edit = True
+        # Disconnect all attrs
+        destination_conns = cmds.listConnections(self, plugs=True, connections=True, source=False) or []
+        for i in range(0, len(destination_conns), 2):
+            cmds.disconnectAttr(destination_conns[i], destination_conns[i + 1])
+        source_conns = cmds.listConnections(self, plugs=True, connections=True, destination=False) or []
+
+        for i in range(0, len(source_conns), 2):
+            # we have to flip these because the output is always node centric and not connection centric
+            cmds.disconnectAttr(source_conns[i + 1], source_conns[i])
+        # Delete the node
+        cmds.delete(self.node)

Scripts/Animation/epic_pose_wrangler/v2/model/utils.py

@@ -0,0 +1,297 @@
+#  Copyright Epic Games, Inc. All Rights Reserved.
+import math
+import traceback
+
+from maya import OpenMaya, cmds
+
+from epic_pose_wrangler.log import LOG
+from epic_pose_wrangler.v2.model import exceptions
+
+# NOTE: MTransformationMatrix & MEulerRotation have different values for the same axis.
+XFORM_ROTATION_ORDER = {
+    'xyz': OpenMaya.MTransformationMatrix.kXYZ,
+    'yzx': OpenMaya.MTransformationMatrix.kYZX,
+    'zxy': OpenMaya.MTransformationMatrix.kZXY,
+    'xzy': OpenMaya.MTransformationMatrix.kXZY,
+    'yxz': OpenMaya.MTransformationMatrix.kYXZ,
+    'zyx': OpenMaya.MTransformationMatrix.kZYX
+}
+
+EULER_ROTATION_ORDER = {
+    'xyz': OpenMaya.MEulerRotation.kXYZ,
+    'yzx': OpenMaya.MEulerRotation.kYZX,
+    'zxy': OpenMaya.MEulerRotation.kZXY,
+    'xzy': OpenMaya.MEulerRotation.kXZY,
+    'yxz': OpenMaya.MEulerRotation.kYXZ,
+    'zyx': OpenMaya.MEulerRotation.kZYX
+}
+
+def compose_matrix(position, rotation, rotation_order='xyz'):
+    """
+    Compose a 4x4 matrix with given transformation.
+
+    >>> compose_matrix((0.0, 0.0, 0.0), (90.0, 0.0, 0.0))
+    [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]
+    """
+
+    # create rotation ptr
+    rot_script_util = OpenMaya.MScriptUtil()
+    rot_script_util.createFromDouble(*[deg * math.pi / 180.0 for deg in rotation])
+    rot_double_ptr = rot_script_util.asDoublePtr()
+
+    # construct transformation matrix
+    xform_matrix = OpenMaya.MTransformationMatrix()
+    xform_matrix.setTranslation(OpenMaya.MVector(*position), OpenMaya.MSpace.kTransform)
+    xform_matrix.setRotation(rot_double_ptr, XFORM_ROTATION_ORDER[rotation_order], OpenMaya.MSpace.kTransform)
+
+    matrix = xform_matrix.asMatrix()
+    return [matrix(m, n) for m in range(4) for n in range(4)]
+
+def decompose_matrix(matrix, rotation_order='xyz'):
+    """
+    Decomposes a 4x4 matrix into translation and rotation.
+
+    >>> decompose_matrix([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0])
+    ((0.0, 0.0, 0.0), (90.0, 0.0, 0.0))
+    """
+
+    mmatrix = OpenMaya.MMatrix()
+    OpenMaya.MScriptUtil.createMatrixFromList(matrix, mmatrix)
+
+    # create transformation matrix
+    xform_matrix = OpenMaya.MTransformationMatrix(mmatrix)
+
+    # get translation
+    translation = xform_matrix.getTranslation(OpenMaya.MSpace.kTransform)
+
+    # get rotation
+    # @ref: https://github.com/LumaPictures/pymel/blob/master/pymel/core/datatypes.py
+    # The apicls getRotation needs a "RotationOrder &" object, which is impossible to make in python...
+    euler_rotation = xform_matrix.eulerRotation()
+    euler_rotation.reorderIt(EULER_ROTATION_ORDER[rotation_order])
+    rotation = euler_rotation.asVector()
+
+    return (
+        (translation.x, translation.y, translation.z),
+        (rotation.x * 180.0 / math.pi, rotation.y * 180.0 / math.pi, rotation.z * 180.0 / math.pi)
+    )
+
+def euler_to_quaternion(rotation, rotation_order='xyz'):
+    """
+    Returns Euler Rotation as Quaternion
+
+    >>> euler_to_quaternion((90, 0, 0))
+    (0.7071, 0.0, 0.0, 0.70710))
+    """
+    euler_rotation = OpenMaya.MEulerRotation(
+        rotation[0] * math.pi / 180.0,
+        rotation[1] * math.pi / 180.0,
+        rotation[2] * math.pi / 180.0
+    )
+    euler_rotation.reorderIt(EULER_ROTATION_ORDER[rotation_order])
+
+    quat = euler_rotation.asQuaternion()
+    return quat.x, quat.y, quat.z, quat.w
+
+def quaternion_to_euler(rotation, rotation_order='xyz'):
+    """
+    Returns Quaternion Rotation as Euler
+
+    quaternion_to_euler((0.7071, 0.0, 0.0, 0.70710))
+    (90, 0, 0)
+    """
+    quat = OpenMaya.MQuaternion(*rotation)
+    euler_rotation = quat.asEulerRotation()
+    euler_rotation.reorderIt(EULER_ROTATION_ORDER[rotation_order])
+
+    return euler_rotation.x * 180.0 / math.pi, euler_rotation.y * 180.0 / math.pi, euler_rotation.z * 180.0 / math.pi
+
+def is_connected_to_array(attribute, array_attr):
+    """
+    Check if the attribute is connected to the specified array
+    :param attribute :type str: attribute
+    :param array_attr :type str array attribute
+    :return :type int or None: int of index in the array or None
+    """
+    try:
+        indices = cmds.getAttr(array_attr, multiIndices=True) or []
+    except ValueError:
+        return None
+    for i in indices:
+        attr = '{from_attr}[{index}]'.format(from_attr=array_attr, index=i)
+        if attribute in (cmds.listConnections(attr, plugs=True) or []):
+            return i
+    return None
+
+def get_next_available_index_in_array(attribute):
+    # Get the next available index
+    indices = cmds.getAttr(attribute, multiIndices=True) or []
+    i = 0
+    for index in indices:
+        if index != i and i not in indices:
+            indices.append(i)
+        i += 1
+    indices.sort()
+    attrs = ['{from_attr}[{index}]'.format(from_attr=attribute, index=i) for i in indices]
+    connections = [cmds.listConnections(attr, plugs=True) or [] for attr in attrs]
+    target_index = len(indices)
+    for index, conn in enumerate(connections):
+        if not conn:
+            target_index = index
+            break
+    return target_index
+
+def message_connect(from_attribute, to_attribute, in_array=False, out_array=False):
+    """
+    Create and connect a message attribute between two nodes
+    """
+    # Generate the object and attr names
+    from_object, from_attribute_name = from_attribute.split('.', 1)
+    to_object, to_attribute_name = to_attribute.split('.', 1)
+
+    # If the attributes don't exist, create them
+    if not cmds.attributeQuery(from_attribute_name, node=from_object, exists=True):
+        cmds.addAttr(from_object, longName=from_attribute_name, attributeType='message', multi=in_array)
+    if not cmds.attributeQuery(to_attribute_name, node=to_object, exists=True):
+        cmds.addAttr(to_object, longName=to_attribute_name, attributeType='message', multi=out_array)
+    # Check that both attributes, if existing are message attributes
+    for a in (from_attribute, to_attribute):
+        if cmds.getAttr(a, type=1) != 'message':
+            raise exceptions.MessageConnectionError(
+                'Message Connect: Attribute {attr} is not a message attribute. CONNECTION ABORTED.'.format(
+                    attr=a
+                )
+            )
+    # Connect up the attributes
+    try:
+        if in_array:
+            from_attribute = "{from_attribute}[{index}]".format(
+                from_attribute=from_attribute,
+                index=get_next_available_index_in_array(from_attribute)
+            )
+
+        if out_array:
+            to_attribute = "{to_attribute}[{index}]".format(
+                to_attribute=to_attribute,
+                index=get_next_available_index_in_array(to_attribute)
+            )
+
+        return cmds.connectAttr(from_attribute, to_attribute, force=True)
+    except Exception as e:
+        LOG.error(traceback.format_exc())
+        return False
+
+def connect_attr(from_attr, to_attr):
+    if not cmds.isConnected(from_attr, to_attr):
+        cmds.connectAttr(from_attr, to_attr)
+
+def get_attr(attr_name, as_value=True):
+    """
+    Get the specified attribute
+    :param attr_name :type str: attribute name i.e node.translate
+    :param as_value :type bool: return as value or connected plug name
+    :return :type list or any: either returns a list of connections or the value of the attribute
+    """
+    # Check if the attribute is connected
+    connections = cmds.listConnections(attr_name, plugs=True)
+    if connections and not as_value:
+        # If the attribute is connected and we don't want the value, return the connections
+        return connections
+    elif as_value:
+        # Return the value
+        return cmds.getAttr(attr_name)
+
+def get_attr_array(attr_name, as_value=True):
+    """
+    Get the specified array attr
+    :param attr_name :type str: attribute name i.e node.translate
+    :param as_value :type bool: return as value or connected plug name
+    :return :type list or any: either returns a list of connections or the value of the attribute
+    """
+    # Get the number of indices in the array
+    indices = cmds.getAttr(attr_name, multiIndices=True) or []
+    # Empty list to store the connected plugs
+    connected_plugs = []
+    # Empty list to store values
+    values = []
+    # Iterate through the indices
+    for i in indices:
+        # Get all the connected plugs for this index
+        connections = cmds.listConnections('{attr_name}[{index}]'.format(attr_name=attr_name, index=i), plugs=True)
+        # If we want the plugs and not values, store connections
+        if connections and not as_value:
+            connected_plugs.extend(connections)
+        # If we want values, get the value at the index
+        elif as_value:
+            values.append(cmds.getAttr('{attr_name}[{index}]'.format(attr_name=attr_name, index=i)))
+    # Return plugs or values, depending on which one has data
+    return connected_plugs or values
+
+def set_attr_or_connect(source_attr_name, value=None, attr_type=None, output=False):
+    """
+    Set an attribute or connect it to another attribute
+    :param source_attr_name :type str: attribute name
+    :param value : type any: value to set the attribute to
+    :param attr_type :type str: name of the attribute type i.e matrix
+    :param output :type bool: is this plug an output (True) or input (False)
+    """
+    # Type conversion from maya: python
+    attr_types = {
+        'matrix': list
+    }
+    # Check if we have a matching type
+    matching_type = attr_types.get(attr_type, None)
+    # If we have a matching type and the value matches that type, set the attr
+    if matching_type is not None and isinstance(value, matching_type):
+        cmds.setAttr(source_attr_name, value, type=attr_type)
+    # If the value is a string and no type is matched, we want to connect the attributes
+    elif isinstance(value, str):
+        try:
+            # Connect from left->right depending on if the source is output or input
+            if output:
+                if not cmds.isConnected(source_attr_name, value):
+                    cmds.connectAttr(source_attr_name, value)
+            else:
+                if not cmds.isConnected(value, source_attr_name):
+                    cmds.connectAttr(value, source_attr_name)
+        except Exception as e:
+            raise exceptions.PoseWranglerAttributeError(
+                "Unable to {direction} {input} to '{output}'".format(
+                    direction="connect" if value else "disconnect",
+                    input=source_attr_name if output else value,
+                    output=value if output else source_attr_name
+                )
+            )
+    else:
+        cmds.setAttr(source_attr_name, value)
+
+def disconnect_attr(attr_name, array=False):
+    """
+    Disconnect the specified attribute
+    :param attr_name :type str: attribute name to disconnect
+    :param array :type bool: is this attribute an array?
+    """
+    attrs = []
+    # If we are disconnecting an array, get the names of all the attributes
+    if array:
+        attrs.extend(cmds.getAttr(attr_name, multiIndices=True) or [])
+    # Otherwise append the attr name specified
+    else:
+        attrs.append(attr_name)
+    # Iterate through all the attrs listed
+    for attr in attrs:
+        # Find their connections and disconnect them
+        for plug in cmds.listConnections(attr, plugs=True) or []:
+            cmds.disconnectAttr(attr, plug)
+
+def get_selection(_type=""):
+    """
+    Returns the current selection
+    """
+    return cmds.ls(selection=True, type=_type)
+
+def set_selection(selection_list):
+    """
+    Sets the active selection
+    """
+    cmds.select(selection_list, replace=True)