Advanced event types

estusshots/views/events.py

@@ -3,11 +3,49 @@ from collections import namedtuple
 from flask import render_template, request, redirect
 
 from estusshots import app, orm
-from estusshots import forms, models, choices
+from estusshots import forms, choices
 from estusshots.util import authorize
 from estusshots.orm import new_session, EventType, Event, Episode, Enemy, Penalty
 
 
+def death_event(event, db, form):
+    """Add penalties to the event if it is a death event"""
+    if not event.id:
+        for entry in form.penalties:
+            penalty = orm.Penalty()
+            penalty.player_id = entry.player_id.data
+            penalty.drink_id = entry.drink.data
+            db.add(penalty)
+            event.penalties.append(penalty)
+    else:
+        for event in event.penalties:
+            penalty = next((p for p in form.penalties.data if p["penalty_id"] == event.id), None)
+            if not penalty:
+                continue
+            penalty.player_id = form.player_id.data
+            penalty.drink_id = form.drink.data
+
+
+def victory_event(event, db, form):
+    """No need for additional actions yet"""
+    pass
+
+
+def pause_event(event, db, form):
+    """Pause events have neither penalties or enemies"""
+    event.enemy_id = None
+    event.enemy = None
+    event.player_id = None
+    event.player = None
+
+
+event_type_map = {
+    EventType.Death: death_event,
+    EventType.Victory: victory_event,
+    EventType.Pause: pause_event
+}
+
+
 @app.route("/season/<s_id>/episode/<ep_id>/event/new", methods=["GET"])
 @authorize
 def event_new(s_id: int, ep_id: int):
@@ -61,21 +99,13 @@ def event_edit(s_id: int, ep_id: int, ev_id: int):
             return render_template("event_editor.html", model=model, form=form)
         if not event:
             event = Event()
-            for entry in form.penalties:
-                penalty = orm.Penalty()
-                penalty.player_id = entry.player_id.data
-                penalty.drink_id = entry.drink.data
-                db.add(penalty)
-                event.penalties.append(penalty)
             db.add(event)
-        else:
-            for event in event.penalties:
-                penalty = next((p for p in form.penalties.data if p["penalty_id"] == event.id), None)
-                if not penalty:
-                    continue
-                penalty.player_id = form.player_id.data
-                penalty.drink_id = form.drink.data
-
         event.populate_from_form(form)
+        # Different event types fill different fields
+        func = event_type_map[event.type]
+        func(event, db, form)
         db.commit()
         return redirect(f"/season/{s_id}/episode/{ep_id}")
+
+
+