main.py 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. import logging
  2. from time import sleep
  3. from telegram import InlineKeyboardMarkup
  4. from telegram.ext import Updater, CallbackQueryHandler, MessageHandler, Filters
  5. from vid_utils import Video, VideoQueue, BadLink
  6. logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
  7. level=logging.INFO)
  8. logger = logging.getLogger(__name__)
  9. VIDEOS = VideoQueue()
  10. def get_format(bot, update):
  11. logger.info("from {}: {}".format(update.message.chat_id, update.message.text)) # "history"
  12. try:
  13. video = Video(update.message.text, update.message.chat_id)
  14. except BadLink:
  15. update.message.reply_text("Bad link")
  16. else:
  17. reply_markup = InlineKeyboardMarkup(video.keyboard)
  18. VIDEOS.append(video)
  19. update.message.reply_text('Choose format:', reply_markup=reply_markup)
  20. def download_choosen_format(bot, update):
  21. query = update.callback_query
  22. bot.edit_message_text(text="Downloading...",
  23. chat_id=query.message.chat_id,
  24. message_id=query.message.message_id)
  25. while VIDEOS.lock: sleep(1) # finish old download
  26. VIDEOS.lock = True # maybe we can use a contextmanager?
  27. for i, video in enumerate(VIDEOS):
  28. if video.chat_id == query.message.chat_id:
  29. VIDEOS.pop(i)
  30. video.download(query.data)
  31. with video.send() as files:
  32. for f in files:
  33. bot.send_document(chat_id=query.message.chat_id, document=open(f, 'rb'))
  34. VIDEOS.lock = False
  35. updater = Updater(token=YOUR_TOKEN)
  36. updater.dispatcher.add_handler(MessageHandler(Filters.text, get_format))
  37. updater.dispatcher.add_handler(CallbackQueryHandler(download_choosen_format))
  38. # Start the Bot
  39. updater.start_polling()
  40. updater.idle()