From a1207d1c119a6fcb98fc64568fd14b43bef7a958 Mon Sep 17 00:00:00 2001 From: blank X Date: Thu, 3 Sep 2020 11:53:54 +0700 Subject: [PATCH] Deal with big images --- redditbot.py | 27 +++++++++++++++++++-------- 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/redditbot.py b/redditbot.py index 7bcda5e..2b72491 100644 --- a/redditbot.py +++ b/redditbot.py @@ -126,17 +126,21 @@ async def main(): break file.write(chunk) + async def _get_file_mimetype(filename): + mimetype = mimetypes.guess_type(filename)[0] + if not mimetype: + proc = await asyncio.create_subprocess_exec('file', '--brief', '--mime-type', filename, stdout=asyncio.subprocess.PIPE) + stdout, _ = await proc.communicate() + mimetype = stdout.decode().strip() + return mimetype or '' + async def _get_file_ext(filename): proc = await asyncio.create_subprocess_exec('file', '--brief', '--extension', filename, stdout=asyncio.subprocess.PIPE) stdout, _ = await proc.communicate() ext = stdout.decode().strip().split('/', maxsplit=1)[0] if not ext or ext == '???': - mimetype = mimetypes.guess_type(filename)[0] - if not mimetype: - proc = await asyncio.create_subprocess_exec('file', '--brief', '--mime-type', filename, stdout=asyncio.subprocess.PIPE) - stdout, _ = await proc.communicate() - mimetype = stdout.decode().strip() - ext = mimetypes.guess_extension(mimetype or '') or '.bin' + mimetype = await _get_file_mimetype(filename) + ext = mimetypes.guess_extension(mimetype) or '.bin' if not ext.startswith('.'): ext = '.' + ext return ext @@ -225,10 +229,17 @@ async def main(): if i in preview: url = preview[i]['source']['url'] break - elif random_post.is_reddit_media_domain and preview and preview['enabled'] and not random_post.is_video: - url = preview['images'][0]['source']['url'] if url: await _download_file(filename, url) + mimetype = await _get_file_mimetype(filename) + if mimetype.startswith('image') and preview and preview['enabled']: + preview = preview['images'][0] + urls = [preview['source']['url']] + urls.extend(reversed(i['url'] for i in preview['resolutions'])) + for url in urls: + if os.path.getsize(filename) < 10000000: + break + await _download_file(filename, url) ext = await _get_file_ext(filename) if ext.startswith('.htm'): files = []