Skip to content

Commit c60f184

Browse files
authored
Merge pull request #84 from Thromax/master
Voting system (!poll command)
2 parents 1cb0b5d + 2dd8179 commit c60f184

File tree

4 files changed

+395
-18
lines changed

4 files changed

+395
-18
lines changed

brainbot.ini.example

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,8 @@
22
tell_me_to=200
33
topic=100
44
repeat=45
5-
phon=45
5+
phon=45
6+
poll=100
7+
8+
[misc]
9+
poll_reactions=zero;one;two;three;four;five;six;seven;eight;nine;keycap_ten

main.py

Lines changed: 269 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
from asyncio import get_event_loop
22
from configparser import ConfigParser
3+
from datetime import datetime, timedelta
34
from os import getenv, system
45
from random import choice
5-
from string import punctuation
66
from sys import executable
77
from urllib.parse import quote
88

@@ -12,10 +12,21 @@
1212
from phonetic_alphabet import read as phonetics
1313
from phonetic_alphabet.main import NonSupportedTextException
1414
from py_expression_eval import Parser
15-
from pyryver import Ryver
16-
from pyryver.util import retry_until_available
17-
18-
from utils import Cooldown, TopicGenerator, bot_dir, console, send_message
15+
from pyryver import Ryver, RyverWS
16+
from pyryver.objects import Notification, TaskBoard
17+
from pyryver.util import datetime_to_iso8601, retry_until_available
18+
from pyryver.ws_data import WSEventData
19+
from pytz import timezone
20+
21+
from utils import (
22+
Cooldown,
23+
TopicGenerator,
24+
bot_dir,
25+
console,
26+
handle_notification,
27+
remind_task,
28+
send_message,
29+
)
1930

2031
__version__ = "1.3.0"
2132

@@ -30,11 +41,20 @@
3041
topic_cooldown = Cooldown(config.getint("cooldowns", "topic", fallback=100))
3142
repeat_cooldown = Cooldown(config.getint("cooldowns", "repeat", fallback=45))
3243
phon_cooldown = Cooldown(config.getint("cooldowns", "phon", fallback=45))
44+
poll_cooldown = Cooldown(config.getint("cooldowns", "poll", fallback=100))
45+
46+
# Load poll reactions
47+
poll_reactions = config.get(
48+
"misc",
49+
"poll_reactions",
50+
fallback="zero;one;two;three;four;five;six;seven;eight;nine;keycap_ten",
51+
).split(";")
3352

3453
math_parser = Parser()
3554
topic_engine = TopicGenerator()
3655
translator = Translator()
3756

57+
3858
# Wrap in async function to use async context manager
3959
async def main():
4060
# Log into Ryver with regular username/password
@@ -56,6 +76,31 @@ async def main():
5676
ryver.get_user(username=user) for user in getenv("BOT_ADMIN").split(",")
5777
]
5878

79+
# Get bot user for task and timezone consults
80+
bot_user = ryver.get_user(id=(await ryver.get_info())["me"]["id"])
81+
82+
# Get bot task board (used for setting timers)
83+
bot_task_board = await bot_user.get_task_board()
84+
if bot_task_board is None:
85+
console.log("Creating task board")
86+
await bot_user.create_task_board(
87+
board_type=TaskBoard.BOARD_TYPE_BOARD, categories=["BrainBot:Polls"]
88+
)
89+
bot_task_board = await bot_user.get_task_board()
90+
elif bot_task_board.get_board_type() == TaskBoard.BOARD_TYPE_LIST:
91+
console.log("Task list in use, BrainBot will not use categories for tasks")
92+
console.log(
93+
"Loaded user {0} task {1}".format(
94+
getenv("RYVER_USER"), bot_task_board.get_board_type()
95+
)
96+
)
97+
98+
# Handle unread notifications from last session (used for checking reminders)
99+
async for notification in ryver.get_notifs(unread=True):
100+
await handle_notification(
101+
ryver=ryver, notification=notification, bot_chat=bot_chat
102+
)
103+
59104
async with ryver.get_live_session() as session:
60105
console.log("In live session")
61106

@@ -226,17 +271,214 @@ async def _on_chat(msg):
226271
# Random Emoticon
227272
elif msg.text.lower().startswith("!emoticon"):
228273
emoticons = [
229-
"( ͡❛ ͜ʖ ͡❛)",
230-
"O_o",
231-
"( ´_ゝ`)",
232-
"(╯°□°)╯︵ ┻━┻",
233-
":-)",
234-
"<(o_o<)",
235-
"(/^▽^)/",
236-
"(•‿•)",
274+
"`( ͡❛ ͜ʖ ͡❛)`",
275+
"`O_o`",
276+
"`( 0ゝ0 )`",
277+
"`(╯°□°)╯︵ ┻━┻`",
278+
"`:-)`",
279+
"`<(o_o<)`",
280+
"`(/^▽^)/`",
281+
"`〠_〠`",
282+
"`(¬‿¬ )`",
283+
"`ᕕ( ᐛ )ᕗ`",
237284
]
238285
console.log(f"Giving {user.get_username()} a random emoticon.")
239286
await send_message(choice(emoticons), bot_chat)
287+
# Create Poll
288+
elif msg.text.lower().startswith("!poll"):
289+
"""
290+
Command usage: !poll [t=due_time;]<poll_title>;<option1>;<option2>...
291+
!poll [d=due_date;]<poll_title>;<option1>;<option2>...
292+
!poll [m=minutes;]<poll_title>;<option1>;<option2>...
293+
[] = Optional <> = Mandatory
294+
Time should use the following format: t=HH:MM
295+
d=mm/dd/yyyy HH:MM
296+
m=<minutes>
297+
Time options d= ,t= and d= are not compatible.
298+
In case of using more than one of them only the first one will be used, while the other
299+
will be considered the rest of the arguments.
300+
Bot timezone will be used.
301+
302+
Poll maximum option number depends of the amount of reactions in the config file
303+
setting 'misc:poll_reactions'.
304+
305+
If due date/time is entered a task and a reminder will be created inside bot's personal
306+
task board in order to make ryver take care of timers instead of the bot itself.
307+
"""
308+
if poll_cooldown.run(username=user.get_username()):
309+
# Get potential arguments
310+
inputs = [value.strip() for value in msg.text[6:].split(";")]
311+
312+
# Remove any empty arguments
313+
while "" in inputs:
314+
inputs.remove("")
315+
316+
# Check if the command contains due date argument
317+
due_date = None
318+
if inputs[0].startswith("t="):
319+
try:
320+
# Parse ending time
321+
due_date = inputs[0][2:]
322+
due_date = datetime.strptime(due_date, "%H:%M")
323+
# Get current time at bot timezone
324+
current_date = datetime.now(
325+
timezone(bot_user.get_time_zone())
326+
)
327+
328+
# If entered hour is earlier (or equal) than the current time, set date due for the next day
329+
if current_date.time() >= due_date.time():
330+
due_date = datetime.combine(
331+
current_date.today() + timedelta(days=1),
332+
due_date.time(),
333+
)
334+
else:
335+
due_date = datetime.combine(
336+
current_date.today(),
337+
due_date.time(),
338+
)
339+
340+
# Set date's timezone
341+
due_date = due_date.astimezone(
342+
timezone(bot_user.get_time_zone())
343+
)
344+
inputs.pop(0)
345+
except ValueError:
346+
due_date = False
347+
348+
if inputs[0].startswith("d="):
349+
try:
350+
# Parse full ending date
351+
due_date = inputs[0][2:]
352+
due_date = datetime.strptime(due_date, "%m/%d/%Y %H:%M")
353+
due_date = due_date.astimezone(
354+
timezone(bot_user.get_time_zone())
355+
)
356+
inputs.pop(0)
357+
except ValueError:
358+
due_date = False
359+
360+
if inputs[0].startswith("m="):
361+
try:
362+
due_date = int(inputs[0][2:])
363+
# Get current time at bot timezone
364+
current_date = datetime.now(
365+
timezone(bot_user.get_time_zone())
366+
)
367+
due_date = current_date.replace(
368+
microsecond=0
369+
) + timedelta(minutes=due_date, seconds=1)
370+
inputs.pop(0)
371+
except ValueError:
372+
due_date = False
373+
374+
if due_date is None or due_date is not False:
375+
current_date = datetime.now(
376+
timezone(bot_user.get_time_zone())
377+
)
378+
379+
# In case of valid due time, check if it's already in the past
380+
if (
381+
due_date is None
382+
or int((due_date - current_date).total_seconds() / 60)
383+
> 0
384+
):
385+
386+
# Check if the command contains a valid number of arguments
387+
if len(inputs) < 3:
388+
await send_message(
389+
"Please enter a question and at least two options to create a poll",
390+
bot_chat,
391+
)
392+
elif len(inputs) > (len(poll_reactions) + 1):
393+
await send_message(
394+
f"Your poll contained too many options, limit is {len(poll_reactions)} options",
395+
bot_chat,
396+
)
397+
else:
398+
console.log(
399+
f'Creating poll "{inputs[0]}" for {user.get_username()}'
400+
)
401+
402+
# Create formatted poll text
403+
poll_txt = "# {0}\n".format(inputs[0])
404+
for i in range(1, len(inputs)):
405+
poll_txt += ":{0}: {1}\n".format(
406+
poll_reactions[i - 1], inputs[i]
407+
)
408+
409+
if due_date is not None:
410+
poll_txt += "\n\n**Poll will end on {0} at {1} ({2})**".format(
411+
due_date.date(),
412+
due_date.time(),
413+
due_date.tzname(),
414+
)
415+
poll_id = await send_message(
416+
poll_txt,
417+
bot_chat,
418+
f"This poll was created by {user.get_username()}",
419+
)
420+
421+
# Get the poll message
422+
message = await retry_until_available(
423+
bot_chat.get_message,
424+
poll_id,
425+
timeout=5.0,
426+
retry_delay=0.5,
427+
)
428+
429+
# Add reaction options
430+
for i in range(0, (len(inputs)) - 1):
431+
await message.react(poll_reactions[i])
432+
433+
# Set ending timer using tasks
434+
if due_date is not None:
435+
task_body = "{0}".format(inputs[0])
436+
# Add options to task message for later parsing
437+
for i in inputs[1:]:
438+
task_body += ";{0}".format(i)
439+
# Add reactions used for later parsing
440+
task_body += ";"
441+
for i in poll_reactions[: (len(inputs) - 1)]:
442+
task_body += ";{0}".format(i)
443+
poll_task = await bot_task_board.create_task(
444+
f"BrainBotPoll#{poll_id}",
445+
task_body,
446+
due_date=datetime_to_iso8601(due_date),
447+
)
448+
# Create task reminder
449+
await remind_task(
450+
ryver,
451+
poll_task,
452+
int(
453+
(
454+
due_date - current_date
455+
).total_seconds()
456+
/ 60
457+
),
458+
)
459+
else:
460+
await send_message(
461+
"Ending time entered is already in the past or too short",
462+
bot_chat,
463+
)
464+
else:
465+
await send_message(
466+
"Ending time entered is not valid. You can any of these formats:\n `t=hh:mm;`\n `d=mm/dd/yyyy hh:mm;`\n`m=<minutes>;`\n**~Don't~ ~forget~ ~to~ ~use~ ~';'!~**",
467+
bot_chat,
468+
)
469+
else:
470+
console.log("Cancelled due to cooldown")
471+
472+
# Get the invoking message
473+
message = await retry_until_available(
474+
bot_chat.get_message,
475+
msg.message_id,
476+
timeout=5.0,
477+
retry_delay=0.5,
478+
)
479+
480+
# React to show the command is on cooldown
481+
await message.react("timer_clock")
240482
# Give a list of commands
241483
elif msg.text.lower().startswith("!commands"):
242484
console.log(f"Telling {user.get_username()} my commands")
@@ -260,7 +502,8 @@ async def _on_chat(msg):
260502
# Render LaTeX
261503
elif msg.text.lower().startswith("!latex"):
262504
await send_message(
263-
f"![LaTeX](http://tex.z-dn.net/?f={quote(msg.text[7:])})", bot_chat
505+
f"![LaTeX](http://tex.z-dn.net/?f={quote(msg.text[7:])})",
506+
bot_chat,
264507
)
265508
# Restart the bot
266509
elif msg.text.lower().startswith("!restart"):
@@ -282,6 +525,18 @@ async def _on_chat(msg):
282525
f"[bold red]{user.get_username()} attempted to shut down the bot"
283526
)
284527

528+
@session.on_event(RyverWS.EVENT_ALL)
529+
async def _on_event(event: WSEventData):
530+
# Check if it's an incoming notification event
531+
# (Notification event type constant doesn't exist on PyRyver, so I hardcoded it)
532+
if event.event_type == "/api/notify":
533+
notif = await Notification.get_by_id(
534+
ryver, obj_id=event.event_data.get("id")
535+
)
536+
await handle_notification(
537+
ryver=ryver, notification=notif, bot_chat=bot_chat
538+
)
539+
285540
@session.on_connection_loss
286541
async def _on_connection_loss():
287542
await session.close()

requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,4 +4,5 @@ phonetic-alphabet==0.1.0
44
py-expression-eval==0.3.10
55
pyryver==0.3.2.post1
66
python-dotenv==0.15.0
7+
pytz==2020.1
78
rich==9.5.1

0 commit comments

Comments
 (0)