From 46c3bbff3cda16f2575a41b59244bdb76a6cf850 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Wed, 3 Apr 2019 15:11:21 +0300
Subject: [PATCH] Simplify voting in polls

---
 mautrix_telegram/commands/telegram/misc.py | 22 ++++++++++++++++++--
 mautrix_telegram/portal.py                 | 24 +++++++++++++---------
 2 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/mautrix_telegram/commands/telegram/misc.py b/mautrix_telegram/commands/telegram/misc.py
index 5a6848c6..7e6299c8 100644
--- a/mautrix_telegram/commands/telegram/misc.py
+++ b/mautrix_telegram/commands/telegram/misc.py
@@ -241,7 +241,7 @@ async def play(evt: CommandEvent) -> Optional[Dict]:
                  help_text="Vote in a Telegram poll.")
 async def vote(evt: CommandEvent) -> Optional[Dict]:
     if len(evt.args) < 2:
-        return await evt.reply("**Usage:** `$cmdprefix+sp vote <poll ID> <choice ID>`")
+        return await evt.reply("**Usage:** `$cmdprefix+sp vote <poll ID> <choice number>`")
     elif not await evt.sender.is_logged_in():
         return await evt.reply("You must be logged in with a real account to vote in polls.")
     elif evt.sender.is_bot:
@@ -255,7 +255,25 @@ async def vote(evt: CommandEvent) -> Optional[Dict]:
     if not isinstance(msg.media, MessageMediaPoll):
         return await evt.reply("Invalid poll ID (message doesn't look like a poll)")
 
-    options = [base64.b64decode(option + (3 - (len(option) + 3) % 4) * "=")
+    options = []
+    for option in evt.args[1:]:
+        try:
+            if len(option) > 10:
+                raise ValueError("option index too long")
+            option_index = int(option) - 1
+        except ValueError:
+            option_index = None
+        if not option_index:
+            return await evt.reply(f"Invalid option number \"{option}\"",
+                                   render_markdown=False, allow_html=False)
+        elif option_index < 0:
+            return await evt.reply(f"Invalid option number {option}. "
+                                   f"Option numbers must be positive.")
+        elif option_index >= len(msg.media.poll.answers):
+            return await evt.reply(f"Invalid option number {option}. "
+                                   f"The poll only has {len(msg.media.poll.answers)} options.")
+        options.append(msg.media.poll.answers[option_index].option)
+    options = [msg.media.poll.answers[int(option) - 1].option
                for option in evt.args[1:]]
     try:
         resp = await evt.sender.client(SendVoteRequest(peer=peer, msg_id=msg.id, options=options))
diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py
index c5368b38..d3a883b2 100644
--- a/mautrix_telegram/portal.py
+++ b/mautrix_telegram/portal.py
@@ -1512,20 +1512,24 @@ class Portal:
         poll = evt.media.poll  # type: Poll
         poll_id = self._encode_msgid(source, evt)
 
-        def enc(answer: PollAnswer) -> str:
-            return base64.b64encode(answer.option).decode("utf-8").rstrip("=")
+        _n = 0
 
-        text = (f"Poll ID {poll_id}: {poll.question}\n"
-                + "\n".join(f"* {enc(answer)}: {answer.text}" for answer in poll.answers) +
+        def n() -> int:
+            nonlocal _n
+            _n += 1
+            return _n
+
+        text = (f"Poll: {poll.question}\n"
+                + "\n".join(f"{n()}. {answer.text}" for answer in poll.answers) +
                 "\n"
-                f"Vote with !tg vote <poll ID> <choice ID>")
+                f"Vote with !tg vote {poll_id} <choice number>")
 
-        html = (f"<strong>Poll</strong> ID <code>{poll_id}</code>: {poll.question}<br/>\n"
-                f"<ul>"
-                + "\n".join(f"<li><code>{enc(answer)}</code>: {answer.text}</li>"
+        html = (f"<strong>Poll</strong>: {poll.question}<br/>\n"
+                f"<ol>"
+                + "\n".join(f"<li>{answer.text}</li>"
                             for answer in poll.answers) +
-                "</ul>\n"
-                f"Vote with <code>!tg vote &lt;poll ID&gt; &lt;choice ID&gt;</code>")
+                "</ol>\n"
+                f"Vote with <code>!tg vote {poll_id} &lt;choice number&gt;</code>")
         await intent.set_typing(self.mxid, is_typing=False)
         return await intent.send_text(self.mxid, text, html=html, relates_to=relates_to,
                                       msgtype="m.text", timestamp=evt.date,