main.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  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. for i, v in enumerate(VIDEOS):
  18. if v.chat_id == video.chat_id:
  19. VIDEOS[i] = video # remove old video not downloaded...
  20. break
  21. else:
  22. VIDEOS.append(video)
  23. reply_markup = InlineKeyboardMarkup(video.keyboard)
  24. update.message.reply_text('Choose format:', reply_markup=reply_markup)
  25. def download_choosen_format(bot, update):
  26. query = update.callback_query
  27. bot.edit_message_text(text="Downloading...",
  28. chat_id=query.message.chat_id,
  29. message_id=query.message.message_id)
  30. while VIDEOS.lock: sleep(1) # finish old download
  31. VIDEOS.lock = True # maybe we can use a contextmanager?
  32. for i, video in enumerate(VIDEOS):
  33. if video.chat_id == query.message.chat_id:
  34. VIDEOS.pop(i)
  35. video.download(query.data)
  36. with video.send() as files:
  37. for f in files:
  38. bot.send_document(chat_id=query.message.chat_id, document=open(f, 'rb'))
  39. VIDEOS.lock = False
  40. updater = Updater(token=YOUR_TOKEN)
  41. updater.dispatcher.add_handler(MessageHandler(Filters.text, get_format))
  42. updater.dispatcher.add_handler(CallbackQueryHandler(download_choosen_format))
  43. # Start the Bot
  44. updater.start_polling()
  45. updater.idle()