Python decorator with optional argument (which is function) -

note: know decorators optional argument contain 3 nested function. optional argument here function itself. please go through complete post before mark duplicate. tried tricks decorators optional argument, not found takes function argument.

i having decorator wrapping error:

def wrap_error(func):     functools import wraps      @wraps(func)     def wrapper(*args, **kwargs):         try:             return func(*args, **kwargs)         except:             import sys              exc_msg = traceback.format_exception(*sys.exc_info())             raise mycustomerror(exc_msg)      return wrapper 

if function raises exception, wraps error. wrapper used like:

@wrap_error def foo():     ... 

now want modify wrapper additional callback function optional. , want wrapper used as:

@wrap_error def foo():     ...  @wrap_error(callback) def foo():     ... 

i know how write decorators optional arguments (in case passed argument not function, based on isfunction(func) check within wrapper). not sure how handle case.

note: can not use @wrap_error() instead of @wrap_error. wrapper used in multiple number of packages, , not possible update change in all

here blocker: consider wrapper as:

@wrap_error(callback)               --->       foo = wrap_error(callback)(foo) def foo():     ... 

so, time wrap_error(foo) executed, not know whether there callback function execution after or not (in case use @wrap_error instead of @wrap_error(callback)).

if there no (callback), wrapping function within wrap_error return func(*args. **kwargs) can raise exception. else have return func called @ next step, , if func() raises exception, call callback() in except block.

since hard tell decorator(func) decorator(callback), make 2 decorators:

from functools import wraps  class mycustomerror(exception):     def __init__(self):         print('in mycustomerror')  # common implementation def wrap(func,cb=none):     @wraps(func)     def wrapper(*args, **kwargs):         try:             return func(*args, **kwargs)         except:             if cb not none:                 cb()             raise mycustomerror()     return wrapper  # no parameters version def wrap_error(func):     return wrap(func)  # callback parameter version def wrap_error_cb(cb):     def deco(func):         return wrap(func,cb)     return deco  @wrap_error def foo(a,b):     print('in foo',a,b)     raise exception('foo exception')  def callback():     print('in callback')  @wrap_error_cb(callback) def bar(a):     print('in bar',a)     raise exception('bar exception') 

check foo , bar correctly using functools.wraps:

>>> foo <function foo @ 0x0000000003f00400> >>> bar <function bar @ 0x0000000003f00598> 

check wrapped functions work:

>>> foo(1,2) in foo 1 2 in mycustomerror traceback (most recent call last):   file "<interactive input>", line 1, in <module>   file "c:\test.py", line 16, in wrapper     raise mycustomerror() mycustomerror >>> bar(3) in bar 3 in callback in mycustomerror traceback (most recent call last):   file "<interactive input>", line 1, in <module>   file "c:\test.py", line 16, in wrapper     raise mycustomerror() mycustomerror 

updated

here's way syntax requested, think above answer clearer.

from functools import wraps  class mycustomerror(exception):     def __init__(self):         print('in mycustomerror')  # common implementation def wrap(func,cb=none):     @wraps(func)     def wrapper(*args, **kwargs):         try:             return func(*args, **kwargs)         except:             if cb not none:                 cb()             raise mycustomerror()     return wrapper  def wrap_error(func_or_cb):     # if function tagged wrap_error_callback     # return decorator returns wrapped function     # callback.     if hasattr(func_or_cb,'cb'):         def deco(func):             return wrap(func,func_or_cb)         return deco     # otherwise, return wrapped function without callback.     return wrap(func_or_cb)  # decorator tag callbacks wrap_error can distinguish them # *regular* functions. def wrap_error_callback(func):     func.cb = true     return func  ### examples of use  @wrap_error def foo(a,b):     print('in foo',a,b)     raise exception('foo exception')  @wrap_error_callback def callback():     print('in callback')  @wrap_error(callback) def bar(a):     print('in bar',a)     raise exception('bar exception')