Added context manager support

This commit is contained in:
Andrew Murray 2017-10-07 23:18:23 +11:00
parent 1d18a5cfef
commit 97662294f0
4 changed files with 24 additions and 12 deletions

View File

@ -420,6 +420,12 @@ class Parser(object):
self.image = im
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
def close(self):
"""
(Consumer) Close the stream.

View File

@ -118,6 +118,12 @@ class ChunkStream(object):
return cid, pos, length
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
def close(self):
self.queue = self.crc = self.fp = None

View File

@ -504,15 +504,15 @@ class TestFilePng(PillowTestCase):
chunks = []
with open(test_file, "rb") as fp:
fp.read(8)
png = PngImagePlugin.PngStream(fp)
while True:
cid, pos, length = png.read()
chunks.append(cid)
try:
s = png.call(cid, pos, length)
except EOFError:
break
png.crc(cid, s)
with PngImagePlugin.PngStream(fp) as png:
while True:
cid, pos, length = png.read()
chunks.append(cid)
try:
s = png.call(cid, pos, length)
except EOFError:
break
png.crc(cid, s)
# https://www.w3.org/TR/PNG/#5ChunkOrdering
# IHDR - shall be first

View File

@ -74,9 +74,9 @@ class TestImageFile(PillowTestCase):
def test_ico(self):
with open('Tests/images/python.ico', 'rb') as f:
data = f.read()
p = ImageFile.Parser()
p.feed(data)
self.assertEqual((48, 48), p.image.size)
with ImageFile.Parser() as p:
p.feed(data)
self.assertEqual((48, 48), p.image.size)
def test_safeblock(self):
if "zip_encoder" not in codecs: