test_imghdr.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. import imghdr
  2. import io
  3. import sys
  4. import unittest
  5. from test.test_support import findfile, TESTFN, unlink, run_unittest
  6. TEST_FILES = (
  7. ('python.png', 'png'),
  8. ('python.gif', 'gif'),
  9. ('python.bmp', 'bmp'),
  10. ('python.ppm', 'ppm'),
  11. ('python.pgm', 'pgm'),
  12. ('python.pbm', 'pbm'),
  13. ('python.jpg', 'jpeg'),
  14. ('python.ras', 'rast'),
  15. ('python.sgi', 'rgb'),
  16. ('python.tiff', 'tiff'),
  17. ('python.xbm', 'xbm')
  18. )
  19. class UnseekableIO(io.FileIO):
  20. def tell(self):
  21. raise io.UnsupportedOperation
  22. def seek(self, *args, **kwargs):
  23. raise io.UnsupportedOperation
  24. class TestImghdr(unittest.TestCase):
  25. @classmethod
  26. def setUpClass(cls):
  27. cls.testfile = findfile('python.png', subdir='imghdrdata')
  28. with open(cls.testfile, 'rb') as stream:
  29. cls.testdata = stream.read()
  30. def tearDown(self):
  31. unlink(TESTFN)
  32. def test_data(self):
  33. for filename, expected in TEST_FILES:
  34. filename = findfile(filename, subdir='imghdrdata')
  35. self.assertEqual(imghdr.what(filename), expected)
  36. ufilename = filename.decode(sys.getfilesystemencoding())
  37. self.assertEqual(imghdr.what(ufilename), expected)
  38. with open(filename, 'rb') as stream:
  39. self.assertEqual(imghdr.what(stream), expected)
  40. with open(filename, 'rb') as stream:
  41. data = stream.read()
  42. self.assertEqual(imghdr.what(None, data), expected)
  43. def test_register_test(self):
  44. def test_jumbo(h, file):
  45. if h.startswith(b'eggs'):
  46. return 'ham'
  47. imghdr.tests.append(test_jumbo)
  48. self.addCleanup(imghdr.tests.pop)
  49. self.assertEqual(imghdr.what(None, b'eggs'), 'ham')
  50. def test_file_pos(self):
  51. with open(TESTFN, 'wb') as stream:
  52. stream.write(b'ababagalamaga')
  53. pos = stream.tell()
  54. stream.write(self.testdata)
  55. with open(TESTFN, 'rb') as stream:
  56. stream.seek(pos)
  57. self.assertEqual(imghdr.what(stream), 'png')
  58. self.assertEqual(stream.tell(), pos)
  59. def test_bad_args(self):
  60. with self.assertRaises(TypeError):
  61. imghdr.what()
  62. with self.assertRaises(AttributeError):
  63. imghdr.what(None)
  64. with self.assertRaises(TypeError):
  65. imghdr.what(self.testfile, 1)
  66. with open(self.testfile, 'rb') as f:
  67. with self.assertRaises(AttributeError):
  68. imghdr.what(f.fileno())
  69. def test_invalid_headers(self):
  70. for header in (b'\211PN\r\n',
  71. b'\001\331',
  72. b'\x59\xA6',
  73. b'cutecat',
  74. b'000000JFI',
  75. b'GIF80'):
  76. self.assertIsNone(imghdr.what(None, header))
  77. def test_missing_file(self):
  78. with self.assertRaises(IOError):
  79. imghdr.what('missing')
  80. def test_closed_file(self):
  81. stream = open(self.testfile, 'rb')
  82. stream.close()
  83. with self.assertRaises(ValueError) as cm:
  84. imghdr.what(stream)
  85. stream = io.BytesIO(self.testdata)
  86. stream.close()
  87. with self.assertRaises(ValueError) as cm:
  88. imghdr.what(stream)
  89. def test_unseekable(self):
  90. with open(TESTFN, 'wb') as stream:
  91. stream.write(self.testdata)
  92. with UnseekableIO(TESTFN, 'rb') as stream:
  93. with self.assertRaises(io.UnsupportedOperation):
  94. imghdr.what(stream)
  95. def test_output_stream(self):
  96. with open(TESTFN, 'wb') as stream:
  97. stream.write(self.testdata)
  98. stream.seek(0)
  99. with self.assertRaises(IOError) as cm:
  100. imghdr.what(stream)
  101. def test_main():
  102. run_unittest(TestImghdr)
  103. if __name__ == '__main__':
  104. test_main()