浏览代码

Add poll hearts option

Kirk Trombley 1 年之前
父节点
当前提交
1197abc796
共有 2 个文件被更改,包括 12 次插入6 次删除
  1. 1 0
      commands/commands/__init__.py
  2. 11 6
      commands/commands/poll.py

+ 1 - 0
commands/commands/__init__.py

@@ -39,6 +39,7 @@ config = get_command_config().extend(
             "wl": "watchlist",
             "choose": "choice",
             "omens": "omen",
+            "pollhearts": "poll",
         },
     )
 )

+ 11 - 6
commands/commands/poll.py

@@ -1,17 +1,18 @@
 import re
 
-from rollbot import as_command, Attachment
+from rollbot import as_command, Attachment, Command
 from rollbot.injection import Args, ArgList, SenderName, Origin
 
 # RGYBbPO
-EMOJIS = ["🟥", "🟩", "🟨", "🟦", "🟫", "🟪", "🟧", "❤️", "💚", "💛", "💙", "🤎", "💜", "🧡"]
+SQUARES = ["🟥", "🟩", "🟨", "🟦", "🟫", "🟪", "🟧"]
+HEARTS = ["❤️", "💚", "💛", "💙", "🤎", "💜", "🧡"]
 
 # Title best effort search
 TITLE_PAT = re.compile(r"\s*[tT]itle:\s*(.+)")
 
 
 @as_command
-def poll(origin: Origin, name: SenderName, args: Args, arg_list: ArgList):
+def poll(origin: Origin, name: SenderName, args: Args, arg_list: ArgList, cmd: Command):
     if origin != "DISCORD":
         return "Polls only supported in Discord!"
 
@@ -29,8 +30,12 @@ def poll(origin: Origin, name: SenderName, args: Args, arg_list: ArgList):
     if title:
         del opts[i]
 
-    if len(opts) > len(EMOJIS):
-        return f"Polls are limited to {len(EMOJIS)} options! Hardware limitation :\\("
+    emojis = (
+        HEARTS + SQUARES if cmd.raw_name.lower() == "pollhearts" else SQUARES + HEARTS
+    )
+
+    if len(opts) > len(emojis):
+        return f"Polls are limited to {len(emojis)} options! Hardware limitation :\\("
 
     if len(opts) == 0:
         return "Can't make an empty poll!"
@@ -39,7 +44,7 @@ def poll(origin: Origin, name: SenderName, args: Args, arg_list: ArgList):
     if title:
         msg += f"**{title}**\n"
     atts = []
-    for e, o in zip(EMOJIS, opts):
+    for e, o in zip(emojis, opts):
         msg += f"> {e} - **{o.strip()}**\n"
         atts.append(Attachment(name="react", body=e))
     msg += "\nPlease participate in democracy by picking the emojis corresponding to your choices!"