1 files changed, 34 insertions, 17 deletions

diff --git a/newspipe/web/views/category.py b/newspipe/web/views/category.py
index cf52b3dd..392b8df6 100644
--- a/newspipe/web/views/category.py
+++ b/newspipe/web/views/category.py
@@ -33,39 +33,30 @@ def list_():
 def form(category_id=None):
     action = gettext("Add a category")
     head_titles = [action]
+    form = CategoryForm()
+    form.set_feed_choices(FeedController(current_user.id).read())
     if category_id is None:
         return render_template(
             "edit_category.html",
             action=action,
             head_titles=head_titles,
-            form=CategoryForm(),
+            form=form,
         )
     category = CategoryController(current_user.id).get(id=category_id)
     action = gettext("Edit category")
     head_titles = [action]
     if category.name:
         head_titles.append(category.name)
+    form = CategoryForm(obj=category)
+    form.set_feed_choices(FeedController(current_user.id).read())
+    form.feeds.data = [feed.id for feed in category.feeds]
     return render_template(
         "edit_category.html",
         action=action,
         head_titles=head_titles,
         category=category,
-        form=CategoryForm(obj=category),
-    )
-
-
-@category_bp.route("/delete/<int:category_id>", methods=["GET"])
-@login_required
-def delete(category_id=None):
-    category = CategoryController(current_user.id).delete(category_id)
-    flash(
-        gettext(
-            "Category %(category_name)s successfully deleted.",
-            category_name=category.name,
-        ),
-        "success",
+        form=form,
     )
-    return redirect(redirect_url())
 
 
 @category_bp.route("/create", methods=["POST"])
@@ -73,7 +64,9 @@ def delete(category_id=None):
 @login_required
 def process_form(category_id=None):
     form = CategoryForm()
+    form.set_feed_choices(FeedController(current_user.id).read())
     cat_contr = CategoryController(current_user.id)
+    feed_contr = FeedController(current_user.id)
 
     if not form.validate():
         return render_template("edit_category.html", form=form)
@@ -81,11 +74,18 @@ def process_form(category_id=None):
     if existing_cats and category_id is None:
         flash(gettext("Couldn't add category: already exists."), "warning")
         return redirect(url_for("category.form", category_id=existing_cats[0].id))
+
     # Edit an existing category
     category_attr = {"name": form.name.data}
-
     if category_id is not None:
         cat_contr.update({"id": category_id}, category_attr)
+        # Update the relation with feeds
+        for feed_id in form.feeds.data:
+            feed_contr.update({"id": feed_id}, {"category_id": category_id})
+        for feed in current_user.feeds:
+            if feed.category_id == category_id and feed.id not in form.feeds.data:
+                feed_contr.update({"id": feed.id}, {"category_id": None})
+
         flash(
             gettext(
                 "Category %(cat_name)r successfully updated.",
@@ -97,6 +97,9 @@ def process_form(category_id=None):
 
     # Create a new category
     new_category = cat_contr.create(**category_attr)
+    # Update the relation with feeds
+    for feed_id in form.feeds.data:
+        feed_contr.update({"id": feed_id}, {"category_id": new_category.id})
 
     flash(
         gettext(
@@ -107,3 +110,17 @@ def process_form(category_id=None):
     )
 
     return redirect(url_for("category.form", category_id=new_category.id))
+
+
+@category_bp.route("/delete/<int:category_id>", methods=["GET"])
+@login_required
+def delete(category_id=None):
+    category = CategoryController(current_user.id).delete(category_id)
+    flash(
+        gettext(
+            "Category %(category_name)s successfully deleted.",
+            category_name=category.name,
+        ),
+        "success",
+    )
+    return redirect(redirect_url())