Use decorators for stall and dish validation

canteeneo/views.py

 import os.path
+from functools import wraps
 from canteeneo import app, login_manager, db
 from flask import flash, redirect, render_template, request, url_for
 from flask_login import login_required, login_user, logout_user, current_user
@@ -51,6 +52,30 @@ def register():
 
     return render_template('register.html', form=form)
 
+def stall_validate(f):
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        stall_id = int(kwargs.get('stall_id'))
+        stall = Stall.query.get(stall_id)
+        if stall is None or stall.owner != current_user:
+            return redirect(url_for('stalls'))
+        return f(*args, **kwargs)
+    return wrapper
+
+def dish_validate(f):
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        stall_id = int(kwargs.get('stall_id'))
+        dish_id = int(kwargs.get('dish_id'))
+        stall = Stall.query.get(stall_id)
+        if stall is None or stall.owner != current_user:
+            return redirect(url_for('stalls'))
+        dish = Dish.query.get(dish_id)
+        if dish is None or dish.stall != stall:
+            return redirect(url_for('stall', stall_id=stall_id))
+        return f(*args, **kwargs)
+    return wrapper
+
 @app.route('/stalls')
 @login_required
 def stalls():
@@ -58,11 +83,9 @@ def stalls():
 
 @app.route('/stalls/<int:stall_id>')
 @login_required
+@stall_validate
 def stall(stall_id):
-    stall = Stall.query.filter_by(id=int(stall_id)).first()
-    if stall is None or stall.owner != current_user:
-        return redirect(url_for('stalls'))
-        
+    stall = Stall.query.get(stall_id)
     dishes = stall.dishes.all()
     return render_template('stall.html', stall=stall, dishes=dishes, upload_folder=app.config['UPLOAD_FOLDER'])
 
@@ -85,11 +108,9 @@ def new_stall():
 
 @app.route('/stalls/<int:stall_id>/edit', methods=['GET', 'POST'])
 @login_required
+@stall_validate
 def edit_stall(stall_id):
-    stall = Stall.query.filter_by(id=int(stall_id)).first()
-    if stall is None or stall.owner != current_user:
-        return redirect(url_to('stalls'))
-
+    stall = Stall.query.get(stall_id)
     form = StallRegisterForm()
     form.location.choices = [(loc.id, loc.name) for loc in Location.query.all()]
     
@@ -112,12 +133,9 @@ def edit_stall(stall_id):
 
 @app.route('/stalls/<int:stall_id>/delete', methods=['POST'])
 @login_required
+@stall_validate
 def delete_stall(stall_id):
-    stall = Stall.query.filter_by(id=int(stall_id)).first()
-
-    if stall is None or stall.owner != current_user:
-        return redirect(url_for('stalls'))
-
+    stall = Stall.query.get(stall_id)
     db.session.delete(stall)
     db.session.commit()
 
@@ -125,11 +143,9 @@ def delete_stall(stall_id):
 
 @app.route('/stalls/<int:stall_id>/dish/new', methods=['GET','POST'])
 @login_required
+@stall_validate
 def new_dish(stall_id):
-    stall = Stall.query.filter_by(id=int(stall_id)).first()
-    if stall is None or stall.owner != current_user:
-        return redirect(url_for('stalls'))
-
+    stall = Stall.query.get(stall_id)
     form = DishRegisterForm()
     if request.method == 'POST':
         if form.validate():
@@ -147,15 +163,10 @@ def new_dish(stall_id):
 
 @app.route('/stalls/<int:stall_id>/dish/<int:dish_id>/edit', methods=['GET', 'POST'])
 @login_required
+@dish_validate
 def edit_dish(stall_id, dish_id):
-    stall = Stall.query.filter_by(id=int(stall_id)).first()
-    if stall is None or stall.owner != current_user:
-        return redirect(url_for('stalls'))  
-    
-    dish = Dish.query.filter_by(id=dish_id).first()
-    if dish is None or dish.stall != stall:
-        return redirect(url_for('stall', stall_id=stall_id))
-
+    stall = Stall.query.get(stall_id)
+    dish = Dish.query.get(dish_id)
     form = DishRegisterForm()
     if request.method == 'POST':
         if form.validate(editing=True):
@@ -179,15 +190,9 @@ def edit_dish(stall_id, dish_id):
 
 @app.route('/stalls/<int:stall_id>/dish/<int:dish_id>/delete', methods=['POST'])
 @login_required
+@dish_validate
 def delete_dish(stall_id, dish_id):
-    stall = Stall.query.filter_by(id=int(stall_id)).first()
-    if stall is None or stall.owner != current_user:
-        return redirect(url_for('stalls'))  
-    
-    dish = Dish.query.filter_by(id=dish_id).first()
-    if dish is None or dish.stall != stall:
-        return redirect(url_for('stall', stall_id=stall_id))
-
+    dish = Dish.query.get(dish_id)
     db.session.delete(dish)
     db.session.commit()