Fix image getting logic

pull/26/head
Guillermo Rodriguez 2023-10-05 16:18:04 -03:00
parent a9b05b7ee7
commit 9fa5a01c34
No known key found for this signature in database
GPG Key ID: EC3AB66D5D83B392
1 changed files with 47 additions and 20 deletions

View File

@ -5,7 +5,7 @@ import random
import logging import logging
import asyncio import asyncio
from PIL import Image from PIL import Image, UnidentifiedImageError
from json import JSONDecodeError from json import JSONDecodeError
from decimal import Decimal from decimal import Decimal
from hashlib import sha256 from hashlib import sha256
@ -244,37 +244,64 @@ class SkynetTelegramFrontend:
async def get_and_set_results(link: str): async def get_and_set_results(link: str):
results[link] = await get_ipfs_file(link) results[link] = await get_ipfs_file(link)
def get_image_from_resp(resp):
png_img = resp.raw
with Image.open(io.BytesIO(resp.raw)) as image:
w, h = image.size
if w > TG_MAX_WIDTH or h > TG_MAX_HEIGHT:
logging.warning(f'result is of size {image.size}')
image.thumbnail((TG_MAX_WIDTH, TG_MAX_HEIGHT))
tmp_buf = io.BytesIO()
image.save(tmp_buf, format='PNG')
png_img = tmp_buf.getvalue()
return png_img
tasks = [ tasks = [
get_and_set_results(ipfs_link), get_and_set_results(ipfs_link),
get_and_set_results(ipfs_link_legacy) get_and_set_results(ipfs_link_legacy)
] ]
await asyncio.gather(*tasks) await asyncio.gather(*tasks)
png_img = None
resp = results[ipfs_link_legacy] resp = results[ipfs_link_legacy]
if not resp or resp.status_code != 200: if not resp or resp.status_code != 200:
logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link_legacy}!') logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link_legacy}!')
resp = results[ipfs_link] else:
if not resp or resp.status_code != 200: try:
logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link}!') png_img = get_image_from_resp(resp)
await self.update_status_message(
status_msg,
caption,
reply_markup=build_redo_menu(),
parse_mode='HTML'
)
return True
png_img = resp.raw except UnidentifiedImageError:
with Image.open(io.BytesIO(resp.raw)) as image: logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link_legacy}!')
w, h = image.size
if w > TG_MAX_WIDTH or h > TG_MAX_HEIGHT: if not png_img:
logging.warning(f'result is of size {image.size}') resp = results[ipfs_link]
image.thumbnail((TG_MAX_WIDTH, TG_MAX_HEIGHT)) if not resp or resp.status_code != 200:
tmp_buf = io.BytesIO() logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link}!')
image.save(tmp_buf, format='PNG') await self.update_status_message(
png_img = tmp_buf.getvalue() status_msg,
caption,
reply_markup=build_redo_menu(),
parse_mode='HTML'
)
return True
else:
try:
png_img = get_image_from_resp(resp)
except UnidentifiedImageError:
logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link}!')
await self.update_status_message(
status_msg,
caption,
reply_markup=build_redo_menu(),
parse_mode='HTML'
)
return True
logging.info(f'success! sending generated image') logging.info(f'success! sending generated image')
await self.bot.delete_message( await self.bot.delete_message(