diff --git a/robot.py b/robot.py index fc02a64..6f0f08c 100644 --- a/robot.py +++ b/robot.py @@ -29,18 +29,21 @@ def date(): return datetime.now().strftime("%Y-%m-%d@%H:%M") -def download_ydl(url): +def download_ydl(urls): ydl_opts = { 'noplaylist': True, 'download_archive': DIR + '/downloaded.lst', 'outtmpl': OUT_DIR + '/' + date() + '__%(title)s__%(id)s.%(ext)s' } with youtube_dl.YoutubeDL(ydl_opts) as ydl: - ydl.download([url]) + ydl.download(urls) def download_raw(url): - local_filename = OUT_DIR + '/' + date() + '__' + url.split('/')[-1] + ext = '' + if 'api.telegram.org' in url and 'animation' in url: + ext = '.mp4' + local_filename = OUT_DIR + '/' + date() + '__' + url.split('/')[-1] + ext r = requests.get(url, stream=True) with open(local_filename, 'wb') as f: for chunk in r.iter_content(chunk_size=1024): @@ -49,12 +52,14 @@ def download_raw(url): def handle_url(bot, update): - for entity in update.message.entities: - if entity.type == 'url': - url = update.message.text[entity.offset:entity.offset + entity.length] - if ytdl_has(url): - update.message.reply_text('Downloading now...') - download_ydl(url) + ytdl_urls = list(filter(ytdl_has, + map(lambda e: update.message.text[e.offset:e.offset + e.length], + filter(lambda e: e.type == 'url', + update.message.entities)))) + if len(ytdl_urls) > 0: + logger.info("Downloading %s" % ytdl_urls) + update.message.reply_text('Downloading now...') + download_ydl(ytdl_urls) def handle_photo(bot, update):