How can I use a Python decorator with arguments?

General Discussions
1

How can I use a Python decorator with arguments?

I am facing an issue with passing the variable insurance_mode to a decorator. I tried the following approach:

@execute_complete_reservation(True)
def test_booking_gta_object(self):
    self.test_select_gta_object()

But this does not work. Is there a better way to solve this problem?

Here is the decorator I am using:

def execute_complete_reservation(test_case, insurance_mode):
    def inner_function(self, *args, **kwargs):
        self.test_create_qsf_query()
        test_case(self, *args, **kwargs)
        self.test_select_room_option()
        if insurance_mode:
            self.test_accept_insurance_crosseling()
        else:
            self.test_decline_insurance_crosseling()
        self.test_configure_pax_details()
        self.test_configure_payer_details()

    return inner_function

How can I correctly implement a Python decorator with arguments in this scenario?

2

You can create a decorator factory that accepts arguments and returns the actual decorator. This allows you to pass parameters, such as insurance_mode, to the decorator.

def execute_complete_reservation(insurance_mode):
    def decorator(test_case):
        def inner_function(self, *args, **kwargs):
            self.test_create_qsf_query()
            test_case(self, *args, **kwargs)
            self.test_select_room_option()
            if insurance_mode:
                self.test_accept_insurance_crosseling()
            else:
                self.test_decline_insurance_crosseling()
            self.test_configure_pax_details()
            self.test_configure_payer_details()
        return inner_function
    return decorator

@execute_complete_reservation(True)
def test_booking_gta_object(self):
    self.test_select_gta_object()

Here, execute_complete_reservation now returns a decorator that can access the insurance_mode argument correctly.

3

If the decorator is used as a standalone, you can pass the arguments directly to the decorator. Here’s an example of how you could modify the original decorator:

def execute_complete_reservation(insurance_mode):
    def decorator(test_case):
        def inner_function(self, *args, **kwargs):
            self.test_create_qsf_query()
            test_case(self, *args, **kwargs)
            self.test_select_room_option()
            if insurance_mode:
                self.test_accept_insurance_crosseling()
            else:
                self.test_decline_insurance_crosseling()
            self.test_configure_pax_details()
            self.test_configure_payer_details()
        return inner_function
    return decorator

@execute_complete_reservation(True)
def test_booking_gta_object(self):
    self.test_select_gta_object()
4

Use functools.wraps for better compatibility with other decorators If you’re combining multiple decorators and want to ensure the wrapped function has the correct signature and behavior, use functools.wraps inside the decorator to preserve the original function’s metadata:

import functools

def execute_complete_reservation(insurance_mode):
    def decorator(test_case):
        @functools.wraps(test_case)
        def inner_function(self, *args, **kwargs):
            self.test_create_qsf_query()
            test_case(self, *args, **kwargs)
            self.test_select_room_option()
            if insurance_mode:
                self.test_accept_insurance_crosseling()
            else:
                self.test_decline_insurance_crosseling()
            self.test_configure_pax_details()
            self.test_configure_payer_details()
        return inner_function
    return decorator

@execute_complete_reservation(True)
def test_booking_gta_object(self):
    self.test_select_gta_object()

This ensures that the original function signature, docstrings, and other metadata are retained while the decorator is applied.