test_imghdr.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. import io
  2. import os
  3. import pathlib
  4. import unittest
  5. import warnings
  6. from test.support import findfile, warnings_helper
  7. from test.support.os_helper import TESTFN, unlink
  8. imghdr = warnings_helper.import_deprecated("imghdr")
  9. TEST_FILES = (
  10. ('python.png', 'png'),
  11. ('python.gif', 'gif'),
  12. ('python.bmp', 'bmp'),
  13. ('python.ppm', 'ppm'),
  14. ('python.pgm', 'pgm'),
  15. ('python.pbm', 'pbm'),
  16. ('python.jpg', 'jpeg'),
  17. ('python-raw.jpg', 'jpeg'), # raw JPEG without JFIF/EXIF markers
  18. ('python.ras', 'rast'),
  19. ('python.sgi', 'rgb'),
  20. ('python.tiff', 'tiff'),
  21. ('python.xbm', 'xbm'),
  22. ('python.webp', 'webp'),
  23. ('python.exr', 'exr'),
  24. )
  25. class UnseekableIO(io.FileIO):
  26. def tell(self):
  27. raise io.UnsupportedOperation
  28. def seek(self, *args, **kwargs):
  29. raise io.UnsupportedOperation
  30. class TestImghdr(unittest.TestCase):
  31. @classmethod
  32. def setUpClass(cls):
  33. cls.testfile = findfile('python.png', subdir='imghdrdata')
  34. with open(cls.testfile, 'rb') as stream:
  35. cls.testdata = stream.read()
  36. def tearDown(self):
  37. unlink(TESTFN)
  38. def test_data(self):
  39. for filename, expected in TEST_FILES:
  40. filename = findfile(filename, subdir='imghdrdata')
  41. self.assertEqual(imghdr.what(filename), expected)
  42. with open(filename, 'rb') as stream:
  43. self.assertEqual(imghdr.what(stream), expected)
  44. with open(filename, 'rb') as stream:
  45. data = stream.read()
  46. self.assertEqual(imghdr.what(None, data), expected)
  47. self.assertEqual(imghdr.what(None, bytearray(data)), expected)
  48. def test_pathlike_filename(self):
  49. for filename, expected in TEST_FILES:
  50. with self.subTest(filename=filename):
  51. filename = findfile(filename, subdir='imghdrdata')
  52. self.assertEqual(imghdr.what(pathlib.Path(filename)), expected)
  53. def test_register_test(self):
  54. def test_jumbo(h, file):
  55. if h.startswith(b'eggs'):
  56. return 'ham'
  57. imghdr.tests.append(test_jumbo)
  58. self.addCleanup(imghdr.tests.pop)
  59. self.assertEqual(imghdr.what(None, b'eggs'), 'ham')
  60. def test_file_pos(self):
  61. with open(TESTFN, 'wb') as stream:
  62. stream.write(b'ababagalamaga')
  63. pos = stream.tell()
  64. stream.write(self.testdata)
  65. with open(TESTFN, 'rb') as stream:
  66. stream.seek(pos)
  67. self.assertEqual(imghdr.what(stream), 'png')
  68. self.assertEqual(stream.tell(), pos)
  69. def test_bad_args(self):
  70. with self.assertRaises(TypeError):
  71. imghdr.what()
  72. with self.assertRaises(AttributeError):
  73. imghdr.what(None)
  74. with self.assertRaises(TypeError):
  75. imghdr.what(self.testfile, 1)
  76. with self.assertRaises(AttributeError):
  77. imghdr.what(os.fsencode(self.testfile))
  78. with open(self.testfile, 'rb') as f:
  79. with self.assertRaises(AttributeError):
  80. imghdr.what(f.fileno())
  81. def test_invalid_headers(self):
  82. for header in (b'\211PN\r\n',
  83. b'\001\331',
  84. b'\x59\xA6',
  85. b'cutecat',
  86. b'000000JFI',
  87. b'GIF80'):
  88. self.assertIsNone(imghdr.what(None, header))
  89. def test_string_data(self):
  90. with warnings.catch_warnings():
  91. warnings.simplefilter("ignore", BytesWarning)
  92. for filename, _ in TEST_FILES:
  93. filename = findfile(filename, subdir='imghdrdata')
  94. with open(filename, 'rb') as stream:
  95. data = stream.read().decode('latin1')
  96. with self.assertRaises(TypeError):
  97. imghdr.what(io.StringIO(data))
  98. with self.assertRaises(TypeError):
  99. imghdr.what(None, data)
  100. def test_missing_file(self):
  101. with self.assertRaises(FileNotFoundError):
  102. imghdr.what('missing')
  103. def test_closed_file(self):
  104. stream = open(self.testfile, 'rb')
  105. stream.close()
  106. with self.assertRaises(ValueError) as cm:
  107. imghdr.what(stream)
  108. stream = io.BytesIO(self.testdata)
  109. stream.close()
  110. with self.assertRaises(ValueError) as cm:
  111. imghdr.what(stream)
  112. def test_unseekable(self):
  113. with open(TESTFN, 'wb') as stream:
  114. stream.write(self.testdata)
  115. with UnseekableIO(TESTFN, 'rb') as stream:
  116. with self.assertRaises(io.UnsupportedOperation):
  117. imghdr.what(stream)
  118. def test_output_stream(self):
  119. with open(TESTFN, 'wb') as stream:
  120. stream.write(self.testdata)
  121. stream.seek(0)
  122. with self.assertRaises(OSError) as cm:
  123. imghdr.what(stream)
  124. if __name__ == '__main__':
  125. unittest.main()