# -*- coding: utf-8 -*-
##  Copyright 2019 Trevor van Hoof and Jan Pijpers.
##  Licensed under the Apache License, Version 2.0
##  Downloaded from https://janpijpers.com or https://gumroad.com/janpijpers
##  See the license file attached or on https://www.janpijpers.com/script-licenses/

'''
Name: mayaUndoExample
Description:

two undo examples
one on how to use the undo stack as decorator function

and as inside function
'''

from maya import cmds
import traceback
import pymel.core as pm

def undo( function ):
    '''
        This is a function decorator.
        you can use it by writing @undo one line obove any function.
        before the function gets called an undo chunk is started.
        when the function ends it gets closed.
        Be carefull, if you call the function in itself (recursively) it will break the undo stack.
    :param function:
    :return:
    '''
    def funcCall(*args,**kwargs):

        result = None 
        try:
            ## here we open the undo chunk and we give it the name of the fuction
            cmds.undoInfo( openChunk= True, chunkname = function.__name__ )
            result = function( *args,**kwargs )
        except Exception as e:
            ## If we have an error we will print the stack of the error
            print traceback.format_exc()
            ## we also make sure the maya ui shows an error.
            pm.displayError( "## Error, see script editor: %s"%e )
        finally:
            ## we always need to close the chunk at the end else we corrupt the stack.
            cmds.undoInfo( closeChunk = True )
        return result
    return funcCall 



@undo
def simpleExampleWithDecorator():
    '''
        So here we have the decorator defined above the function definition.
        So before this function is called, an undo chunk is created and closed after the function has finished.
    :return:
    '''
    for i in xrange( 10 ):
        loc = cmds.createNode( "spaceLocator"  )
        cmds.xform(loc, translation=(i,i,i))
simpleExampleWithDecorator()


def undoInFunction( ):

    ## its recomended to always use a try and except with an undo chunk else you need to restart maya when it fails.
    ## here we open the undo chunk and we give it the name of the fuction
    cmds.undoInfo(openChunk=True, chunkname="Example")

    try:
        for i in xrange(10):
            loc = cmds.createNode("spaceLocator")
            cmds.xform(loc, translation=(i, i, i) )
    except Exception as e:
        ## If we have an error we will print the stack of the error
        print traceback.format_exc()
        ## we also make sure the maya ui shows an error.
        pm.displayError("## Error, see script editor: %s" % e)
    finally:
        cmds.undoInfo(openChunk=True, chunkname="Example")

undoInFunction()