2 files changed, 38 insertions, 18 deletions
diff --git a/app/admin/forms.py b/app/admin/forms.py
index 6e48ee9..0c35933 100644
--- a/app/admin/forms.py
+++ b/app/admin/forms.py
@@ -6,8 +6,8 @@ from flask import current_app
 from flask_wtf import FlaskForm
 from flask_wtf.file import FileField, FileRequired
 
-from wtforms import StringField, SubmitField, SelectField
-from wtforms.validators import DataRequired, ValidationError, InputRequired
+from wtforms import SubmitField, SelectField
+from wtforms.validators import ValidationError, InputRequired
 
 from werkzeug.utils import secure_filename
 
@@ -15,12 +15,13 @@ from app.models import Dataset
 from app.admin.datasets import validate_dataset, get_name_from_dataset
 
 
-class AdminAssignTaskForm(FlaskForm):
+class AdminManageTaskForm(FlaskForm):
     username = SelectField(
         "Username", coerce=int, validators=[InputRequired()]
     )
     dataset = SelectField("Dataset", coerce=int, validators=[InputRequired()])
-    submit = SubmitField("Assign")
+    assign = SubmitField("Assign")
+    delete = SubmitField("Delete")
 
 
 class AdminAddDatasetForm(FlaskForm):
diff --git a/app/admin/routes.py b/app/admin/routes.py
index 7617943..bcd4603 100644
--- a/app/admin/routes.py
+++ b/app/admin/routes.py
@@ -10,17 +10,17 @@ from app import db
 from app.admin import bp
 from app.admin.datasets import get_name_from_dataset, md5sum
 from app.admin.decorators import admin_required
-from app.admin.forms import AdminAssignTaskForm, AdminAddDatasetForm
+from app.admin.forms import AdminManageTaskForm, AdminAddDatasetForm
 from app.models import User, Dataset, Task
 
 
-@bp.route("/assign", methods=("GET", "POST"))
+@bp.route("/manage", methods=("GET", "POST"))
 @admin_required
-def assign():
+def manage():
     user_list = [(u.id, u.username) for u in User.query.all()]
     dataset_list = [(d.id, d.name) for d in Dataset.query.all()]
 
-    form = AdminAssignTaskForm()
+    form = AdminManageTaskForm()
     form.username.choices = user_list
     form.dataset.choices = dataset_list
 
@@ -28,26 +28,45 @@ def assign():
         user = User.query.filter_by(id=form.username.data).first()
         if user is None:
             flash("User does not exist.")
-            return redirect(url_for("admin.assign"))
+            return redirect(url_for("admin.manage"))
         dataset = Dataset.query.filter_by(id=form.dataset.data).first()
         if dataset is None:
             flash("Dataset does not exist.")
-            return redirect(url_for("admin.assign"))
+            return redirect(url_for("admin.manage"))
+
+        action = None
+        if form.assign.data:
+            action = "assign"
+        elif form.delete.data:
+            action = "delete"
+        else:
+            flash("Internal error: no button is true but form was submitted.")
+            return redirect(url_for("admin.manage"))
 
         task = Task.query.filter_by(
             annotator_id=user.id, dataset_id=dataset.id
         ).first()
-        if not task is None:
-            flash("Task assignment already exists.")
-            return redirect(url_for("admin.assign"))
+        if task is None:
+            if action == "delete":
+                flash("Can't delete a task that doesn't exist.")
+                return redirect(url_for("admin.manage"))
+            else:
+                task = Task(annotator_id=user.id, dataset_id=dataset.id)
+                db.session.add(task)
+                db.session.commit()
+                flash("Task registered successfully.")
+        else:
+            if action == "assign":
+                flash("Task assignment already exists.")
+                return redirect(url_for("admin.manage"))
+            else:
+                db.session.delete(task)
+                db.session.commit()
+                flash("Task deleted successfully.")
 
-        task = Task(annotator_id=user.id, dataset_id=dataset.id)
-        db.session.add(task)
-        db.session.commit()
-        flash("Task registered successfully.")
     tasks = Task.query.all()
     return render_template(
-        "admin/assign.html", title="Assign Task", form=form, tasks=tasks
+        "admin/manage.html", title="Assign Task", form=form, tasks=tasks
     )