Added context manager, __del__ and close methods

This commit is contained in:
Andrew Murray 2018-11-11 16:00:17 +11:00
parent c501eedcb5
commit 52c175d1bc
2 changed files with 37 additions and 18 deletions

View File

@ -15,21 +15,25 @@ class TestFileTar(PillowTestCase):
self.skipTest("neither jpeg nor zip support available") self.skipTest("neither jpeg nor zip support available")
def test_sanity(self): def test_sanity(self):
if "zip_decoder" in codecs: for codec, test_path, format in [
tar = TarIO.TarIO(TEST_TAR_FILE, 'hopper.png') ['zip_decoder', 'hopper.png', 'PNG'],
['jpeg_decoder', 'hopper.jpg', 'JPEG']
]:
if codec in codecs:
tar = TarIO.TarIO(TEST_TAR_FILE, test_path)
im = Image.open(tar) im = Image.open(tar)
im.load() im.load()
self.assertEqual(im.mode, "RGB") self.assertEqual(im.mode, "RGB")
self.assertEqual(im.size, (128, 128)) self.assertEqual(im.size, (128, 128))
self.assertEqual(im.format, "PNG") self.assertEqual(im.format, format)
if "jpeg_decoder" in codecs: def test_close(self):
tar = TarIO.TarIO(TEST_TAR_FILE, 'hopper.jpg') tar = TarIO.TarIO(TEST_TAR_FILE, 'hopper.jpg')
im = Image.open(tar) tar.close()
im.load()
self.assertEqual(im.mode, "RGB") def test_contextmanager(self):
self.assertEqual(im.size, (128, 128)) with TarIO.TarIO(TEST_TAR_FILE, 'hopper.jpg') as tar:
self.assertEqual(im.format, "JPEG") pass
if __name__ == '__main__': if __name__ == '__main__':

View File

@ -14,6 +14,7 @@
# See the README file for information on usage and redistribution. # See the README file for information on usage and redistribution.
# #
import sys
from . import ContainerIO from . import ContainerIO
@ -30,11 +31,11 @@ class TarIO(ContainerIO.ContainerIO):
:param tarfile: Name of TAR file. :param tarfile: Name of TAR file.
:param file: Name of member file. :param file: Name of member file.
""" """
fh = open(tarfile, "rb") self.fh = open(tarfile, "rb")
while True: while True:
s = fh.read(512) s = self.fh.read(512)
if len(s) != 512: if len(s) != 512:
raise IOError("unexpected end of tar file") raise IOError("unexpected end of tar file")
@ -50,7 +51,21 @@ class TarIO(ContainerIO.ContainerIO):
if file == name: if file == name:
break break
fh.seek((size + 511) & (~511), 1) self.fh.seek((size + 511) & (~511), 1)
# Open region # Open region
ContainerIO.ContainerIO.__init__(self, fh, fh.tell(), size) ContainerIO.ContainerIO.__init__(self, self.fh, self.fh.tell(), size)
# Context Manager Support
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
if sys.version_info.major >= 3:
def __del__(self):
self.close()
def close(self):
self.fh.close()