Search code examples
pythonflaskdecoratorpython-decorators

Flask custom decorator not functioning as expected


I'm using Flask and have two decorators that I'm trying to use on certain routes to be more Pythonic, prevent code reuse, and improve readability.

The first of these are just checks to see if the user is logged in else redirect user to the login page. This works fine.

After this I check for the user 'class' (Admin, Manager, Standard User, etc), however this one isn't working and I'm not sure what is missing.

From the routes.py:

@app.route('/user/user-account.html', methods=['GET', 'POST'])
@login_required
def cna_account():
    if valid_user() == True:
        result = get_user_account()
        return render_template('user/user-account.html', result=result)
    else:
        return redirect(url_for('index'))

@login_reqiured works fine, here is the code:

@wraps(f)
def wrap(*args, **kwargs):
    if 'logged_in' in session:
        return f(*args, **kwargs)
    else:
        return redirect(url_for('index'))
return wrap

What I tried for the other decorator that doesn't work:

def valid_user(f): ''' Makes sure that only Base Users can view the Base User pages '''

@wraps(f)
def wrap(*args, **kwargs):
    if 'access' in session and session['access'] == 'c':
        return f(*args, **kwargs)
    else:
        return redirect(url_for('index'))

These are stored in a separate module that is imported in routes.py, the only thing I can guess is perhaps the session isn't being carried across even though it's included from flask in the module and routes, but again I'm not sure how this would be remedied.

What I'm trying to do with it is: have the routes that use both decorators and don't require the if valid_user() == True check. Instead it should function and look like:

@app.route('/user/user-account.html', methods=['GET', 'POST'])
@login_required
@valid_user
def cna_account():
    result = get_user_account()
    return render_template('user/user-account.html', result=result)

Any help on what I'm missing here? Do I need to pass the session variable as an argument to @valid_user? I tried that a few different ways and still had it throw errors.

Much appreciated!


Solution

  • You need to return the inner function wrap in valid_user:

    def valid_user(f): 
      @wraps(f)
      def wrap(*args, **kwargs):
        if 'access' in session and session['access'] == 'c':
          return f(*args, **kwargs)
        return redirect(url_for('index'))
      return wrap
    

    While it is perfectly valid to not return anything from an outer decorator function, bear in mind that the wrapped function is passed to the wrapper at run time, thus, the returned value will be None:

    def foo(f):
       print("inside decorator with '{}'".format(f.__name__))
       def inner():
          return 10
    
    @foo
    def bar():
      return 'in bar'
    
    "inside decorator with 'bar'"
    >>>bar()
    

    Traceback (most recent call last): File "", line 1, in TypeError: 'NoneType' object is not callable