tale_handler.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  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, update, context):
  20. """Prompt user to tell a tale. It will be added to database if upvoted
  21. at least three times
  22. """
  23. ctype=update.message.chat.type
  24. if update.message.chat.type != Chat.GROUP and update.message.chat.type != Chat.SUPERGROUP :
  25. update.message.reply_text(text='You need to be in the official LCM group chat to 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, update, context):
  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, update, context):
  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. query.answer(text="Nope, can't upvote your own tale")
  66. elif voter.id in t.upvoters_id:
  67. query.answer(text='One upvote per person, sorry')
  68. else:
  69. t.upvoters_id.append(voter.id)
  70. t.upvotes += 1
  71. query.answer(text='Upvote registered')
  72. if t.upvotes == 3:
  73. n_tale = self.db.incr('n_tales')
  74. self.db.hset(str(n_tale), 'author', t.author.first_name)
  75. self.db.hset(str(n_tale), 'tale', t.text)
  76. query.edit_message_text(text='Tale saved to database')
  77. else:
  78. thumb_up = u'\U0001F44D'
  79. button = InlineKeyboardButton(thumb_up + str(t.upvotes),
  80. callback_data=query.data)
  81. keyboard = InlineKeyboardMarkup([[button]])
  82. query.edit_message_text(text=query.message.text,reply_markup=keyboard)
  83. def tell_a_tale(self, update, context):
  84. """Tell a tale about LCM"""
  85. msg = update.message
  86. n_tales = self.db.get('n_tales')
  87. if not n_tales:
  88. text = ("I don't know any tales or anecdotes about LCM yet. "
  89. 'Teach me one using the /addatale command!')
  90. msg.reply_text(text=text)
  91. else:
  92. rnd_tale = str(np.random.randint(1, int(n_tales) + 1))
  93. t = self.db.hgetall(rnd_tale)
  94. msg.reply_text(text="This one's from " + t['author'] + ":",
  95. quote=False)
  96. msg.reply_text(text=t['tale'], quote=False)