fix: create tempfile in the outer scope

Signed-off-by: Frost Ming <me@frostming.com>
This commit is contained in:
Frost Ming 2023-12-12 14:33:50 +08:00
parent 23484f5b7c
commit 82b209fb22
No known key found for this signature in database
GPG Key ID: 5BFA9CB4DDA943BF

View File

@ -58,10 +58,9 @@ def sizeof_image(image):
return f.tell() return f.tell()
def compress_image(input_image, target_size): def compress_image(input_image, output_image, target_size):
quality = 95 quality = 95
factor = 1.0 factor = 1.0
output = SpooledTemporaryFile(max_size=MAX_IN_MEMORY)
with Image.open(input_image) as img: with Image.open(input_image) as img:
while sizeof_image(img) > target_size: while sizeof_image(img) > target_size:
factor -= 0.05 factor -= 0.05
@ -70,20 +69,20 @@ def compress_image(input_image, target_size):
(int(width * factor), int(height * factor)), (int(width * factor), int(height * factor)),
PIL.Image.Resampling.LANCZOS, PIL.Image.Resampling.LANCZOS,
) )
img.save(output, format="JPEG", quality=quality) img.save(output_image, format="JPEG", quality=quality)
output.seek(0) output_image.seek(0)
return output
def draw_pretty_map(location, style): def draw_pretty_map(location, style, output_file):
aoi = get_aoi(address=location, radius=1100, rectangular=True) aoi = get_aoi(address=location, radius=1100, rectangular=True)
df = get_osm_geometries(aoi=aoi) df = get_osm_geometries(aoi=aoi)
fig = Plot(df=df, aoi_bounds=aoi.bounds, draw_settings=STYLES[style]).plot_all() fig = Plot(df=df, aoi_bounds=aoi.bounds, draw_settings=STYLES[style]).plot_all()
with SpooledTemporaryFile(max_size=MAX_IN_MEMORY) as buffer: with SpooledTemporaryFile(max_size=MAX_IN_MEMORY) as buffer:
fig.savefig(buffer, format="jpeg") fig.savefig(buffer, format="jpeg")
buffer.seek(0) buffer.seek(0)
return compress_image( compress_image(
buffer, buffer,
output_file,
10 * 1024 * 1024, # telegram tog need png less than 10MB 10 * 1024 * 1024, # telegram tog need png less than 10MB
) )
@ -151,15 +150,15 @@ def main():
style = random.choice(styles_list) style = random.choice(styles_list)
try: try:
# TODO why this memory leak? # TODO why this memory leak?
out_image = draw_pretty_map(location, style) with SpooledTemporaryFile(max_size=MAX_IN_MEMORY) as out_image:
# tg can only send image less than 10MB draw_pretty_map(location, style, out_image)
with open("map_out.jpg", "wb") as f: # tg can only send image less than 10MB
shutil.copyfileobj(out_image, f) with open("map_out.jpg", "wb") as f: # for debug
out_image.seek(0) shutil.copyfileobj(out_image, f)
bot.send_photo( out_image.seek(0)
message.chat.id, out_image, reply_to_message_id=message.message_id bot.send_photo(
) message.chat.id, out_image, reply_to_message_id=message.message_id
out_image.close() )
except Exception: except Exception:
traceback.print_exc() traceback.print_exc()