Using decorators with argument in Python

Wed Jun 29 17:29:45 EDT 2011

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~