code improvement

This commit is contained in:
Hodza Alban 2017-08-19 18:41:50 +02:00 committed by GitHub
parent 6fd4285b3b
commit d384efb138

36
main.py
View File

@ -1,63 +1,47 @@
import logging import logging
from time import sleep
from telegram import InlineKeyboardMarkup from telegram import InlineKeyboardMarkup
from telegram.ext import Updater, CallbackQueryHandler, MessageHandler, Filters from telegram.ext import Updater, CallbackQueryHandler, MessageHandler, Filters
from vid_utils import Video, VideoQueue, BadLink from vid_utils import Video, BadLink
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)
level=logging.INFO)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
VIDEOS = VideoQueue()
def get_format(bot, update): def get_format(bot, update):
logger.info("from {}: {}".format(update.message.chat_id, update.message.text)) # "history" logger.info("from {}: {}".format(update.message.chat_id, update.message.text)) # "history"
try: try:
video = Video(update.message.text, update.message.chat_id) video = Video(update.message.text, init_keyboard=True)
except BadLink: except BadLink:
update.message.reply_text("Bad link") update.message.reply_text("Bad link")
else: else:
for i, v in enumerate(VIDEOS):
if v.chat_id == video.chat_id:
VIDEOS[i] = video # remove old video not downloaded...
break
else:
VIDEOS.append(video)
reply_markup = InlineKeyboardMarkup(video.keyboard) reply_markup = InlineKeyboardMarkup(video.keyboard)
update.message.reply_text('Choose format:', reply_markup=reply_markup) update.message.reply_text('Choose format:', reply_markup=reply_markup)
def download_choosen_format(bot, update): def download_choosen_format(bot, update):
query = update.callback_query query = update.callback_query
resolution_code, link = query.data.split(' ', 1)
bot.edit_message_text(text="Downloading...", bot.edit_message_text(text="Downloading...",
chat_id=query.message.chat_id, chat_id=query.message.chat_id,
message_id=query.message.message_id) message_id=query.message.message_id)
while VIDEOS.lock: sleep(1) # finish old download video = Video(link)
video.download(resolution_code)
VIDEOS.lock = True # maybe we can use a contextmanager?
for i, video in enumerate(VIDEOS):
if video.chat_id == query.message.chat_id:
VIDEOS.pop(i)
video.download(query.data)
with video.send() as files: with video.send() as files:
for f in files: for f in files:
bot.send_document(chat_id=query.message.chat_id, document=open(f, 'rb')) bot.send_document(chat_id=query.message.chat_id, document=open(f, 'rb'))
VIDEOS.lock = False
updater = Updater(token=YOUR_TOKEN) updater = Updater(token=YOUR_TOKEN)
updater.dispatcher.add_handler(MessageHandler(Filters.text, get_format)) updater.dispatcher.add_handler(MessageHandler(Filters.text, get_format))
updater.dispatcher.add_handler(CallbackQueryHandler(download_choosen_format)) updater.dispatcher.add_handler(CallbackQueryHandler(download_choosen_format))
# Start the Bot
updater.start_polling() updater.start_polling()
updater.idle() updater.idle()