The extends decorator

Recently, I was writing a pipeline tool that spanned over a couple of classes. Each of these classes were actually manipulating Maya nodes. So I've a Sequence class manipulating Sequence nodes in the maya scene and a Shot class manipulating Shot nodes but are in communication with my Sequence class and so on.

After a couple of days of coding I've started to feel that what I was doing was not continuous in terms of the experience you got out of them while coding with them. I mean, you first create your own Sequence instances and let them find and store pymel Sequence instances etc.. Maya Sequence nodes are already wrapped with pymel.core.nt.Sequence class and I'm writing a new class that stores the pymel instance and does what it does on top of it. Although, doing so is the general practice in Maya, but I didn't like it this time.

What I would prefer was to directly be able to manipulate the pymel class to have the methods I've written. Then it flashed suddenly to me. The next one is not what I've came up with, it is a step lead me to my final solution.

So, in Python you can easily patch a class in runtime like this:

def create_shot(self, name='', handle=10):
    """Creates a new shot.

    :param str name: A string value for the newly created shot name, if
      skipped or given empty, the next empty shot name will be generated.
    :param int handle: An integer value for the handle attribute. Default
      is 10.
    :returns: The created :class:`~pymel.core.nt.Shot` instance
    """
    shot = pymel.core.createNode('shot', name=name)
    shot.shotName.set(name)
    self.set_shot_handles([shot], handle=handle)
    # connect to the sequencer
    shot.message >> self.shots.next_available  # this is another story
    return shot

# and patch the Sequence class
pymel.core.nodetypes.Sequence.create_shot = create_shot

Even though, this seems neat/frightening at first sight, I don't like this kind of patching. The patching is good but the way you patch it is not Pythonic enough to my taste.

What I came up with is the following:

def extends(cls):
    """A decorator for extending classes with other class methods or functions.

    :param cls: The class object that will be extended.
    """
    def wrapper(f):
        if isinstance(f, property):
            name = f.fget.__name__
        else:
            name = f.__name__
        setattr(cls, name, f)

        def wrapped_f(*args, **kwargs):
            return f(*args, **kwargs)

        return wrapped_f
    return wrapper

The above decorator, I think, is kind of the most simplest complex code I've ever written. And used neatly as follows:

class SequenceExtension(object):
    """Extension to pymel.core.nt.Sequence class
    """
    @extends(pymel.core.nodetypes.Sequence)
    def create_shot(self, name='', handle=10):
        """Creates a new shot.

        :param str name: A string value for the newly created shot name, if
          skipped or given empty, the next empty shot name will be generated.
        :param int handle: An integer value for the handle attribute. Default
          is 10.
        :returns: The created :class:`~pymel.core.nt.Shot` instance
        """
        shot = pymel.core.createNode('shot', name=name)
        shot.shotName.set(name)
        self.set_shot_handles([shot], handle=handle)
        # connect to the sequencer
        shot.message >> self.shots.next_available
        return shot

So by using it, any Sequence node you get out of Pymel (pm.ls(type=pm.nt.Sequence)) will have that method, event the ones you've already created in runtime. And all of the methods are staying in a good place under their own class. So you can now do things like that:

seqs = pm.ls(type=pm.nt.Sequence)
shot1 = seqs[0].create_shot('test_shot')

Isn't that beauty.