From 4b1144a7577f0bc253f45e91a79bea19f7cd9f0a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Ml=C3=A1dek?= Date: Mon, 5 Feb 2018 09:39:38 +0100 Subject: [PATCH] tweak filename --- robot.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/robot.py b/robot.py index f1c5608..baf6983 100755 --- a/robot.py +++ b/robot.py @@ -69,8 +69,8 @@ def download_ydl(urls, subdir, date, extract=False): shutil.move(globbed, out_dir) -def download_raw(url, subdir, date): - local_filename = f"{OUT_DIR}/{subdir}/" + "%s__%s" % (datestr(date), url.split('/')[-1]) +def download_raw(url, filename, subdir, date): + local_filename = f"{OUT_DIR}/{subdir}/" + "%s__%s" % (datestr(date), filename or url.split('/')[-1]) # local_filename = OUT_DIR + '/' + ("%s/" % subdir) if subdir else '' + datestr(date) + '__' + url.split('/')[-1] r = requests.get(url, stream=True) with open(local_filename, 'wb') as f: @@ -122,13 +122,15 @@ def handle_url(bot, update): # noinspection PyBroadException def handle_rest(bot, update): - file = None + file, filename = None, None if len(update.message.photo) > 0: photo = max(update.message.photo, key=lambda p: p.width) file = photo.file_id elif update.message.document is not None: + filename = update.message.document.file_name file = update.message.document.file_id elif update.message.audio is not None: + filename = update.message.audio.title file = update.message.audio.file_id elif update.message.video is not None: file = update.message.video.file_id @@ -147,7 +149,7 @@ def handle_rest(bot, update): update.message.reply_text(reply) if hashtag: mkdir_p(f'{OUT_DIR}/{hashtag}') - download_raw(url, hashtag or '.', update.message.date) + download_raw(url, filename, hashtag or '.', update.message.date) except: type, value, _ = sys.exc_info() update.message.reply_text("Something is FUCKED: %s, %s" % (type, value)) @@ -173,7 +175,8 @@ def main(): dp.add_error_handler(error) dp.add_handler(MessageHandler(Filters.entity(MessageEntity.URL), handle_url)) - dp.add_handler(MessageHandler(Filters.photo | Filters.video | Filters.audio | Filters.voice | Filters.document, handle_rest)) + dp.add_handler( + MessageHandler(Filters.photo | Filters.video | Filters.audio | Filters.voice | Filters.document, handle_rest)) updater.start_polling()