put functionality on models

Author: cmdelatorre

src/pycamp_bot/models.py

@@ -89,13 +89,24 @@ def set_as_only_active(self):
         self.active = True
         self.save()
 
+    def add_wizard(self, username, chat_id):
+        pycampista = Pycampista.get_or_create(username=username, chat_id=chat_id)[0]
+        pycampista.wizard = True
+        pycampista.save()
+        PycampistaAtPycamp.get_or_create(pycamp=self, pycampista=pycampista)
+        return pycampista
+
     def get_wizards(self):
-        return Pycampista.select().where(Pycampista.wizard == 1)
+        pac = PycampistaAtPycamp.select().join(Pycampista).where(
+            (PycampistaAtPycamp.pycamp == self) &
+            (PycampistaAtPycamp.pycampista.wizard == True)
+        )
+        return [p.pycampista for p in pac]
 
     def get_current_wizard(self):
         """Return the Pycampista instance that's the currently scheduled wizard."""
         now = datetime.now()
-        current_wizards = WizardAtPycamp.select().where(  
+        current_wizards = WizardAtPycamp.select().where(
             (WizardAtPycamp.pycamp == self) & 
             (WizardAtPycamp.init <= now) &
             (WizardAtPycamp.end > now) 

src/pycamp_bot/utils.py

@@ -1,4 +1,6 @@
 from pycamp_bot.models import Pycamp
+from pycamp_bot.logger import logger
+
 
 def escape_markdown(string):
     # See: https://core.telegram.org/bots/api#markdownv2-style
@@ -30,3 +32,21 @@ def get_slot_weekday_name(slot_day_code):
     day_name = ISO_WEEKDAY_NAMES[pycamp_start_weekday + offset]
 
     return day_name
+def active_pycamp_needed(f):
+    from pycamp_bot.commands.manage_pycamp import get_active_pycamp
+    async def wrap(*args, **kargs):
+        update, context = args
+
+        _, pycamp = get_active_pycamp()
+        if pycamp is None:
+            msg = "🔥 %s: This operation (%s) needs an active PyCamp. Talk to an admin." % (
+                update.message.from_user.username, str(f.__name__)
+            )
+            await context.bot.send_message(
+                chat_id=update.message.chat_id,
+                text=msg)
+            logger.warning(msg)
+            return
+
+        return await f(*args, pycamp=pycamp)
+    return wrap