Search code examples
pythonpython-3.xdecoratorpython-decorators

Using wrapper function of decorators to make changes to a function's arguments


I'm having a little trouble wrapping my head around decorators. I want to use the wrapper function inside the decorator to apply changes to my request.Session() object.

This is what the function looks like with my attempt as a decorator:

import os
import requests
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry

# want to make changes to my session object and return the new object
# as an argument to my original function

def decorator_function(func):
    def wrapper(*args, **kwargs):
        session = requests.Session()
        retry = Retry(connect=2, backoff_factor=0.2)
        adapter = HTTPAdapter(max_retries=retry)
        session.mount('http://', adapter)
        session.mount('https://', adapter)
    
        return func(*args, **kwargs)
    return wrapper

To be able to feed a session argument into my wrapper function's arguments, I made session to be keyword argument in my original function:

def download_image(SAVE_DIR, image_links, session=None):
    if not session:
        session = requests.Session()

    for num, link in enumerate(image_links):
        try:
            r = s.get(link)
        except (requests.exceptions.RequestException, UnicodeError) as e:
            print(e)
    
        with open(os.path.join(
                  SAVE_DIR, 'image_{}.jpg'.format(num)),
                  'wb') as f:
            f.write(r.content)

As you can see session has the default value of None, but I'm attempting to use the decorator to create a special session object and feed that into function, overriding the default value of None.

I used **args and **kwargs as arguments to the wrapper function because if I have other functions with additional arguments, I could still use this decorator with it.

However, it seems If I use *args, *kwargs, I can't make changes to the session argument.

Is there an elegant way to use the decorator function to make changes to a session object, and still retain the generality to wrap it around other functions that may have more different arguments in addition to the session object?


Solution

  • Change this line return func(*args, **kwargs)

    To return func(*args, session =session, **kwargs)

    The problem is that you created the session without actually doing anything with it.