Using decorators with argument in Python

[Ethan Furman]

Ian Kelly wrote:
> On Wed, Jun 29, 2011 at 1:30 PM, Ethan Furman <ethan at stoneleaf.us> wrote:
>> How about just having one bit of code that works either way?
> 
> How would you adapt that code if you wanted to be able to decorate a
> function that takes arguments?

8<----------------------------------------------------------------
class enclose(object):
     func = None
     def __init__(self, char='#'):
         self.char = char
         if callable(char):  # was a function passed in directly?
             self.char = '#' # use default char
             self.func = char
     def __call__(self, func=None, *args, **kwargs):
         if self.func is None:
             self.func = func
             return self
         if func is not None:
             args = (func, ) + args
         return self._call(*args, **kwargs)
     def _call(self, *args, **kwargs):
         print("\n" + self.char * 50)
         self.func(*args, **kwargs)
         print(self.char * 50 + '\n')
if __name__ == '__main__':
     @enclose
     def test1():
         print('Spam!')
     @enclose('-')
     def test2():
         print('Eggs!')
     @enclose
     def test3(string):
         print(string)
     @enclose('^')
     def test4(string):
         print(string)
     test1()
     test2()
     test3('Python')
     test4('Rules!  ;)')
8<----------------------------------------------------------------

> This also won't work if the argument to the decorator is itself a
> callable, such as in the OP's example.

Indeed.  In that case you need two keywords to __init__, and the 
discipline to always use the keyword syntax at least for the optional 
function paramater.  On the bright side, if one forgets, it blows up 
pretty quickly.

Whether it's worth the extra effort depends on the programmer's tastes, 
of course.

8<----------------------------------------------------------------
class enclose(object):
     func = None
     pre_func = None
     def __init__(self, dec_func=None, opt_func=None):
         if opt_func is None:
             if dec_func is not None: # was written without ()'s
                 self.func = dec_func
         else:
             self.pre_func = opt_func
     def __call__(self, func=None, *args, **kwargs):
         if self.func is None:
             self.func = func
             return self
         if func is not None:
             args = (func, ) + args
         if self.pre_func is not None:
             self.pre_func()
         return self._call(*args, **kwargs)
     def _call(self, *args, **kwargs):
         print("\n" + '~' * 50)
         self.func(*args, **kwargs)
         print('~' * 50 + '\n')

if __name__ == '__main__':
     def some_func():
         print('some func here')
     @enclose
     def test1():
         print('Spam!')
     @enclose(opt_func=some_func)
     def test2():
         print('Eggs!')
     @enclose
     def test3(string):
         print(string)
     @enclose(opt_func=some_func)
     def test4(string):
         print(string)
     test1()
     test2()
     test3('Python')
     test4('Rules!  ;)')
8<----------------------------------------------------------------

~Ethan~