Replace sys.stdout with sys.stdout.buffer when saving

This commit is contained in:
Andrew Murray 2021-04-25 14:21:00 +10:00
parent ef9a8e5f7f
commit fae9afe351
3 changed files with 22 additions and 1 deletions

View File

@ -1,4 +1,5 @@
import re
import sys
import zlib
from io import BytesIO
@ -10,6 +11,7 @@ from .helper import (
PillowLeakTestCase,
assert_image,
assert_image_equal,
assert_image_equal_tofile,
hopper,
is_big_endian,
is_win32,
@ -711,6 +713,23 @@ class TestFilePng:
with pytest.raises(EOFError):
im.seek(1)
def test_save_stdout(self):
old_stdout = sys.stdout.buffer
class MyStdOut:
buffer = BytesIO()
sys.stdout = mystdout = MyStdOut()
with Image.open(TEST_PNG_FILE) as im:
im.save(sys.stdout, "PNG")
# Reset stdout
sys.stdout = old_stdout
reloaded = Image.open(mystdout.buffer)
assert_image_equal_tofile(reloaded, TEST_PNG_FILE)
@pytest.mark.skipif(is_win32(), reason="Requires Unix or macOS")
@skip_unless_feature("zlib")

View File

@ -2135,6 +2135,8 @@ class Image:
elif isinstance(fp, Path):
filename = str(fp)
open_fp = True
elif fp == sys.stdout:
fp = sys.stdout.buffer
if not filename and hasattr(fp, "name") and isPath(fp.name):
# only set the name for metadata purposes
filename = fp.name

View File

@ -493,7 +493,7 @@ def _save(im, fp, tile, bufsize=0):
# But, it would need at least the image size in most cases. RawEncode is
# a tricky case.
bufsize = max(MAXBLOCK, bufsize, im.size[0] * 4) # see RawEncode.c
if fp == sys.stdout:
if fp == sys.stdout.buffer:
fp.flush()
return
try: