Home > Enterprise >  How to wrap a third party decorator with my own in python
How to wrap a third party decorator with my own in python

Time:12-01

I see a lot of examples of wrapping decorators with your own, but they are all about custom decorators. I have a third party decorator that can be used such as:

@someSDK
myFunction()

what I want to do is wrap this in my own decorator that calls @someSDK only if an env is set to true

EDIT: specifically, the 3rd party decorator is @xray_recorder.capture() . I want my own decorator that will only implement this if I have a variable set as a certain flag.

This is what I have so far as a rough idea:

if CUSTOM_XRAY_WRAPPER is True:
    @xray_recorder.capture()
    def wrapper(*args, **kwargs):
        val = func(*args, **kwargs)
        return val
    return wrapper
else:
    def wrapper(*args, **kwargs):
        val = func(*args, **kwargs)
        return val
    return wrapper

CodePudding user response:

As a little example, since you have not provided a minimal reproducible example:

CUSTOM_XRAY_WRAPPER = True

def xray_capture(func):
    def wrapper(*args, **kwargs):
        print('recodring')
        return func(*args, **kwargs)
    return wrapper

def yourdecorator(func):
    if CUSTOM_XRAY_WRAPPER:
        return xray_capture(func)
    return func

@yourdecorator
def yourfunction(text):
    print(text)

Here, using the xray_capture decorator on a function prints "recording" before executing that function. The function yourdecorator can be used as a decorator instead that returns either the normal function as is if CUSTOM_XRAY_WRAPPER is False. If it's True, it returns the function decorated with xray_capture.

Btw, you don't have to write CUSTOM_XRAY_WRAPPER == True since it is already either True or False, which means you are checking if True == True or False == True.

Transfered to your example, this would look like this:

def yourdecorator(func):
    if CUSTOM_XRAY_WRAPPER:
        xray_recorder.capture()(func)
    return func

CodePudding user response:

What you need to remember is that

@decorator
def foobar():
     ...

is equivalent to

def foobar():
   ...
foobar = decorator(foobar)

So what you want is

def my_decorator(f):
    if CUSTOM_XRAY_WRAPPER:
       return xray_recorder.capture()(f)
    else:
       return f

CodePudding user response:

You could rename the decorator and use a pass-through:

from functools import wraps


MY_ENV_VAR = True


# create a pass-through wrapper
def sample_decorator(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        return f(*args, **kwargs)
    return wrapper


# now you can just use the same decorator everywhere
if MY_ENV_VAR:
    # use your recorder.capture 
    x_ray_capture = xray_recorder.capture
else:
    # pass-through
    x_ray_capture = sample_decorator


@x_ray_capture
def foo():
    # do something
  • Related