get_objects_under_node.py

import sys
import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx

pluginCommandName       = 'getObjsUnderNode'
flagExactNameShort      = '-en'
flagExactNameLong       = '-exactName'
flagFilterNameShort     = '-fn'
flagFilterNameLong      = '-filterName'

class GetObjsUnderNode(OpenMayaMPx.MPxCommand):
    """
        The flags with names are just flags, true/false, 1/2 etc.
        A command argument is something like an object passed in as default to perform behaviour on
        e.g. cmds.myCommand(pCube1, myFlag=True)
    """
    def __init__(self):
        OpenMayaMPx.MPxCommand.__init__(self)

        self.objectPath     = None
        self.exactNameText  = None
        self.filterString   = None
        self.selectionList  = OpenMaya.MSelectionList()

    def doIt(self, args):
        self.parseArguments(args)
        dagObject = self.getDagObjectFromString(self.objectPath)
        self.filteredObjects = self.iterateDagObject(dagObject, filterString=self.filterString)
        self.setResult(self.filteredObjects)

    def parseArguments(self, args):
        # The following MArgParser object allows you to check if specific flags are set.
        argData = OpenMaya.MArgParser( self.syntax(), args )

        if argData.isFlagSet(flagExactNameShort):
            # We print the flags values as string
            # We use the '0' to index the flag's first and only parameter.
            self.exactNameText = argData.flagArgumentString(flagExactNameShort, 0)

        if argData.isFlagSet(flagFilterNameShort):
            # We print the flags values as string
            # We use the '0' to index the flag's first and only parameter.
            self.filterString = argData.flagArgumentString(flagFilterNameShort, 0)

        if argData.commandArgumentString(0):
            self.objectPath = argData.commandArgumentString(0)


    def getDagObjectFromString(self, nodeString):
        dagPath = OpenMaya.MDagPath()
        if nodeString:
            self.selectionList.add(nodeString)
            self.selectionList.getDagPath(0, dagPath)
            dagObject = dagPath.node()
            return dagObject

    def iterateDagObject(self, dagObject, filterString=''):
        iterator = OpenMaya.MItDag()
        iterator.reset(dagObject, OpenMaya.MItDag.kDepthFirst, OpenMaya.MFn.kTransform)

        # Create an MDagPath object which will be populated on each iteration.
        dagPath = OpenMaya.MDagPath()
        # Obtain a reference to MFnDag function set to print the name of the DAG object
        dagFn = OpenMaya.MFnDagNode()

        foundResults = []

        while not iterator.isDone():
            iterator.getPath(dagPath)
            dagObject = dagPath.node()
            dagFn.setObject(dagObject)
            fullPath = dagFn.fullPathName()
            nodeName = dagFn.name()
            if nodeName.startswith(self.filterString):
                foundResults.append(fullPath[1:])
            iterator.next()
        return foundResults

def cmdCreator():
    return OpenMayaMPx.asMPxPtr(GetObjsUnderNode())

def syntaxCreator():
    ''' Defines the argument and flag syntax for this command. '''
    syntax = OpenMaya.MSyntax()

    # We add our flags here and specify type using MSyntax.kString, Boolean, Double etc.
    syntax.addArg(OpenMaya.MSyntax.kString)
    syntax.enableQuery(True)
    syntax.addFlag(flagExactNameShort, flagExactNameLong, OpenMaya.MSyntax.kString)
    syntax.addFlag(flagFilterNameShort, flagFilterNameLong, OpenMaya.MSyntax.kString)
    return syntax
 
def initializePlugin(mObject):
    plugin = OpenMayaMPx.MFnPlugin(mObject, 'Ben Hearn')
    try:
        plugin.registerCommand(pluginCommandName, cmdCreator, syntaxCreator)
    except:
        raise RuntimeError, 'Failed to register command'
        #sys.stderr.write( "Failed to register command: %s\n" % kPluginCmdName )
        #raise
 
def uninitializePlugin(mObject):
    plugin = OpenMayaMPx.MFnPlugin(mObject)
    try:
        plugin.deregisterCommand(pluginCommandName)
    except:
        raise RuntimeError, 'Failed to unregister command'