tale_handler.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119
  1. # -*- coding: utf-8 -*-
  2. import numpy as np
  3. from telegram import (ForceReply, InlineKeyboardMarkup, InlineKeyboardButton,
  4. Chat)
  5. import redis
  6. class Tale:
  7. def __init__(self, text, msg_id, author):
  8. self.text = text
  9. self.msg_id = msg_id
  10. self.author = author
  11. self.upvotes = 0
  12. self.upvoters_id = []
  13. class TaleHandler:
  14. """Helper class to handle storage and retrieval of tales via Telegram"""
  15. def __init__(self):
  16. self.db = redis.StrictRedis()
  17. self.cur_authors = [] # list of tuples (user_id, prompt_msg_id)
  18. self.cur_tales = [] # list of Tale objects (tales to be added to db)
  19. def prompt_user(self, bot, update):
  20. """Prompt user to tell a tale. It will be added to database if upvoted
  21. at least three times
  22. """
  23. if update.message.chat.type != Chat.GROUP:
  24. update.message.reply_text(text='We need to be in a group chat to '
  25. 'add a new tale or anecdote!')
  26. return
  27. user = update.message.from_user
  28. prompt_msg = update.message.reply_text(
  29. text='Ok, %s, tell me your story!' % user.first_name,
  30. reply_markup=ForceReply(selective=True))
  31. self.cur_authors.append((user.id, prompt_msg.message_id))
  32. def handle_new_tale(self, bot, update):
  33. """Start the upvoting procedure when a new tale is received"""
  34. tale_msg = update.message
  35. author = tale_msg.from_user
  36. reply_to_msg = tale_msg.reply_to_message
  37. reply_to_id = reply_to_msg.message_id if reply_to_msg is not None else 0
  38. if (author.id, reply_to_id) not in self.cur_authors:
  39. # this is not a tale. do nothing
  40. return
  41. self.cur_authors.remove((author.id, reply_to_id))
  42. t = Tale(text=tale_msg.text,
  43. msg_id=tale_msg.message_id,
  44. author=author)
  45. self.cur_tales.append(t)
  46. thumb_up = u'\U0001F44D'
  47. b = InlineKeyboardButton(thumb_up + ' 0', callback_data=str(t.msg_id))
  48. keyboard = InlineKeyboardMarkup([[b]])
  49. tale_msg.reply_text(
  50. text='Alright, I need three upvotes to add this tale to my database',
  51. reply_markup=keyboard)
  52. def save_tale(self, bot, update):
  53. """Add tale to db when upvoted three times"""
  54. query = update.callback_query
  55. voter = query.from_user
  56. # retrieve tale
  57. t = filter(lambda t: t.msg_id == int(query.data), self.cur_tales)
  58. if len(t) == 0:
  59. # could not find this tale. this should never happen
  60. # TODO log an error
  61. return
  62. t = t[0]
  63. if voter.id == t.author.id:
  64. # authors cannot upvote their own story
  65. bot.answer_callback_query(
  66. callback_query_id=query.id,
  67. text="Nope, can't upvote your own tale")
  68. elif voter.id in t.upvoters_id:
  69. bot.answer_callback_query(
  70. callback_query_id=query.id,
  71. text='One upvote per person, sorry')
  72. else:
  73. t.upvoters_id.append(voter.id)
  74. t.upvotes += 1
  75. bot.answer_callback_query(
  76. callback_query_id=query.id,
  77. text='Upvote registered')
  78. if t.upvotes == 3:
  79. n_tale = self.db.incr('n_tales')
  80. self.db.hset(str(n_tale), 'author', t.author.first_name)
  81. self.db.hset(str(n_tale), 'tale', t.text)
  82. bot.editMessageText(chat_id=query.message.chat_id,
  83. message_id=query.message.message_id,
  84. text='Tale saved to database')
  85. else:
  86. thumb_up = u'\U0001F44D'
  87. button = InlineKeyboardButton(thumb_up + str(t.upvotes),
  88. callback_data=query.data)
  89. keyboard = InlineKeyboardMarkup([[button]])
  90. bot.editMessageText(chat_id=query.message.chat_id,
  91. message_id=query.message.message_id,
  92. text=query.message.text,
  93. reply_markup=keyboard)
  94. def tell_a_tale(self, bot, update):
  95. """Tell a tale about LCM"""
  96. msg = update.message
  97. n_tales = self.db.get('n_tales')
  98. if not n_tales:
  99. text = ("I don't know any tales or anecdotes about LCM yet. "
  100. 'Teach me one using the /addatale command!')
  101. msg.reply_text(text=text)
  102. else:
  103. rnd_tale = str(np.random.randint(1, int(n_tales) + 1))
  104. t = self.db.hgetall(rnd_tale)
  105. msg.reply_text(text="This one's from " + t['author'] + ":",
  106. quote=False)
  107. msg.reply_text(text=t['tale'], quote=False)