main.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. #!/bin/env python
  2. import logging
  3. import os
  4. import random
  5. from collections import defaultdict, namedtuple
  6. from telegram import ReplyKeyboardMarkup, ReplyKeyboardRemove
  7. from telegram.ext import BaseFilter, CommandHandler, Filters, MessageHandler, Updater
  8. from telegram.utils.helpers import mention_html
  9. from config import TOKEN
  10. logging.basicConfig(
  11. format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
  12. )
  13. class FilterDart(BaseFilter):
  14. def filter(self, message):
  15. return message.dice.emoji == "🎯"
  16. def get_score(value):
  17. return {1: 0, 2: 1, 3: 2, 4: 3, 5: 4, 6: 6}.get(value)
  18. PlayerState = namedtuple("PlayerState", ["score", "throws"])
  19. def format_name(user):
  20. return user.first_name + ((" " + user.last_name) if user.last_name else "")
  21. custom_keyboard = ReplyKeyboardMarkup([["🎯"]])
  22. class ChatGame:
  23. def __init__(self, updater, update, ctx):
  24. self.updater = updater
  25. self.jq = updater.job_queue
  26. self.chat_id = update.effective_chat.id
  27. ctx.bot.send_message(
  28. chat_id=self.chat_id, text="k", reply_markup=custom_keyboard
  29. )
  30. self.states = defaultdict(lambda: PlayerState(0, 0))
  31. print("Starting game in chat {}".format(self.chat_id))
  32. self.replyjobs = dict()
  33. self.replystats = defaultdict(list)
  34. self.namecache = dict()
  35. self.dead = False
  36. def stop(self, update, ctx):
  37. textlines = [
  38. "{}: {} ({} throws) ({:.4f} average)".format(
  39. mention_html(i, self.namecache[i]),
  40. state.score,
  41. state.throws,
  42. state.score / state.throws,
  43. )
  44. for i, state in sorted(
  45. self.states.items(), key=lambda item: item[1].score, reverse=True
  46. )
  47. ]
  48. ctx.bot.send_message(
  49. chat_id=update.effective_chat.id,
  50. parse_mode="HTML",
  51. text="\n".join(textlines),
  52. reply_markup=ReplyKeyboardRemove(),
  53. )
  54. self.states = None
  55. self.dead = True
  56. def dart(self, update, ctx):
  57. dice = update.message.dice
  58. user = update.message.from_user
  59. self.namecache[user.id] = format_name(user)
  60. old = self.states[user.id]
  61. delta = get_score(dice.value)
  62. state = PlayerState(old.score + delta, old.throws + 1)
  63. self.states[user.id] = state
  64. self.replystats[user.id].append(delta)
  65. cnt = len(self.replystats[user.id])
  66. rj = self.replyjobs.get(user.id, None)
  67. if rj:
  68. # XXX maybe race condition?
  69. rj.schedule_removal()
  70. def reply(ctx):
  71. newthrows = self.replystats[user.id][:cnt]
  72. if len(newthrows) > 100:
  73. s = "Not showing individual scores."
  74. else:
  75. s = " ".join((f"+{score}" for score in newthrows))
  76. update.message.reply_text(
  77. "{}\nThrows: {} (+{})\nScore: {} (+{})".format(
  78. s, state.throws, cnt, state.score, sum(newthrows),
  79. )
  80. )
  81. self.replyjobs[user.id] = None
  82. self.replystats[user.id] = self.replystats[user.id][cnt:]
  83. self.replyjobs[user.id] = self.jq.run_once(reply, 2.5)
  84. class DartboiBot:
  85. def __init__(self):
  86. self.updater = Updater(token=TOKEN, use_context=True)
  87. self.dp = self.updater.dispatcher
  88. handlers = [
  89. CommandHandler("start", self.start_command),
  90. CommandHandler("stop", self.meta_handler(ChatGame.stop)),
  91. MessageHandler(
  92. Filters.dice & FilterDart() & (~Filters.forwarded),
  93. self.meta_handler(ChatGame.dart),
  94. ),
  95. ]
  96. for h in handlers:
  97. self.dp.add_handler(h)
  98. self.games = dict()
  99. def run(self):
  100. self.updater.start_polling()
  101. self.updater.idle()
  102. def start_command(self, update, ctx):
  103. chatid = update.effective_chat.id
  104. if chatid in self.games:
  105. update.message.reply_text(random.choice(("あほか?", "ばか!")))
  106. return
  107. game = ChatGame(self.updater, update, ctx)
  108. self.games[chatid] = game
  109. def meta_handler(self, func):
  110. def inner(update, ctx):
  111. chat_id = update.effective_chat.id
  112. game = self.games.get(chat_id, None)
  113. if game is not None:
  114. func(game, update, ctx)
  115. if game.dead:
  116. del self.games[chat_id]
  117. return inner
  118. if __name__ == "__main__":
  119. DartboiBot().run()