diff --git a/PIL/ImageFile.py b/PIL/ImageFile.py index b0d13a0fc..458857f41 100644 --- a/PIL/ImageFile.py +++ b/PIL/ImageFile.py @@ -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. diff --git a/PIL/PngImagePlugin.py b/PIL/PngImagePlugin.py index 409d7f302..454a4d330 100644 --- a/PIL/PngImagePlugin.py +++ b/PIL/PngImagePlugin.py @@ -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 diff --git a/Tests/test_file_png.py b/Tests/test_file_png.py index 200c93b74..644f38e8f 100644 --- a/Tests/test_file_png.py +++ b/Tests/test_file_png.py @@ -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 diff --git a/Tests/test_imagefile.py b/Tests/test_imagefile.py index 6ef74989d..b50dcc943 100644 --- a/Tests/test_imagefile.py +++ b/Tests/test_imagefile.py @@ -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: