From f6a38cc2dbc026e0333f3801e3d77aa386edf4d0 Mon Sep 17 00:00:00 2001
From: luxick <git@luxick.de>
Date: Wed, 13 Mar 2019 20:32:55 +0100
Subject: [PATCH] Refactor code for generating choices lists.

---
 choices.py | 44 +++++++++++++++++++++++++++++++-------------
 forms.py   |  2 +-
 2 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/choices.py b/choices.py
index 972e669..9f16bf6 100644
--- a/choices.py
+++ b/choices.py
@@ -22,23 +22,41 @@ def player_choice():
     return [(p.id, p.name) for p in players]
 
 
-class SeasonChoicesIterable:
+def drink_choice():
+    """
+    Query database for a list of all available drinks to select from
+    """
+    sql, args = db.load_drinks()
+    drinks = db.query_db(sql, args, cls=models.Drink)
+    choices = [(d.id, d.name) for d in drinks]
+    choices.insert(0, (-1, "None"))
+    return choices
+
+
+class IterableBase:
     """
     This is used to declare choices for WTForms SelectFields at class definition time.
     Be aware that this uses an undocumented WTForms feature and is not guaranteed to work.
     """
 
-    def __iter__(self):
-        for choice in season_choices():
-            yield choice
-
-
-class PlayerChoiceIterable:
-    """
-    This is used to declare choices for WTForms SelectFields at class definition time.
-    Be aware that this uses an undocumented WTForms feature and is not guaranteed to work.
-    """
+    _loader = None
 
     def __iter__(self):
-        for choice in player_choice():
-            yield choice
+        if self._loader:
+            for choice in self._loader():
+                yield choice
+
+
+class SeasonChoiceIterable(IterableBase):
+    def __init__(self):
+        self._loader = season_choices
+
+
+class PlayerChoiceIterable(IterableBase):
+    def __init__(self):
+        self._loader = player_choice
+
+
+class DrinkChoiceIterable(IterableBase):
+    def __init__(self):
+        self._loader = drink_choice
diff --git a/forms.py b/forms.py
index 8e5a842..7655aca 100644
--- a/forms.py
+++ b/forms.py
@@ -58,7 +58,7 @@ class DrinkForm(FlaskForm):
 class EnemyForm(FlaskForm):
     enemy_id = HiddenField("Enemy ID", render_kw={"readonly": True})
     season_id = SelectField(
-        "Season", choices=choices.SeasonChoicesIterable(), coerce=int
+        "Season", choices=choices.SeasonChoiceIterable(), coerce=int
     )
     name = StringField("Name", validators=[DataRequired()])
     is_boss = BooleanField("Is Boss")