Rewrite imghdr inlining for clarity and completeness

* Move MIMEImage class back closer to the top of the file since it's the important thing. * Use a decorate to mark a given rule function and simplify the rule function names for clarity. * Copy over all the imghdr test data files into the email package's test data directory. This way when imghdr is actually removed, it won't affect the MIMEImage guessing tests. * Rewrite and extend the MIMEImage tests to test for all supported auto-detected MIME image subtypes. * Remove the now redundant PyBanner048.gif data file.
warsaw · Apr 15, 2022 · Apr 14, 2022 · Apr 14, 2022 · Apr 14, 2022 · Apr 14, 2022
commit e2d4aa6f755e76c49e388f8d830a25e77194b879
@@ -10,137 +10,143 @@
 from email.mime.nonmultipart import MIMENonMultipart
 
 
+class MIMEImage(MIMENonMultipart):
+    """Class for generating image/* type MIME documents."""
+
+    def __init__(self, _imagedata, _subtype=None,
+                 _encoder=encoders.encode_base64, *, policy=None, **_params):
+        """Create an image/* type MIME document.
+
+        _imagedata is a string containing the raw image data.  If the data
+        type can be detected (jpeg, png, gif, tiff, rgb, pbm, pgm, ppm,
+        rast, xbm, bmp, webp, and exr attempted), then the subtype will be
+        automatically included in the Content-Type header. Otherwise, you can
+        specify the specific image subtype via the _subtype parameter.
+
+        _encoder is a function which will perform the actual encoding for
+        transport of the image data.  It takes one argument, which is this
+        Image instance.  It should use get_payload() and set_payload() to
+        change the payload to the encoded form.  It should also add any
+        Content-Transfer-Encoding or other headers to the message as
+        necessary.  The default encoding is Base64.
+
+        Any additional keyword arguments are passed to the base class
+        constructor, which turns them into parameters on the Content-Type
+        header.
+        """
+        if _subtype is None:
+            if (_subtype := _what(_imagedata)) is None:
+                raise TypeError('Could not guess image MIME subtype')
+        MIMENonMultipart.__init__(self, 'image', _subtype, policy=policy,
+                                  **_params)
+        self.set_payload(_imagedata)
+        _encoder(self)
+
+
+_rules = []
+
+
 # Originally from the imghdr module.
-def _what(h):
-    for tf in tests:
-        if res := tf(h):
+def _what(data):
+    for rule in _rules:
+        if res := rule(data):
             return res
     else:
         return None
 
-tests = []
 
-def _test_jpeg(h):
+def rule(rulefunc):
+    _rules.append(rulefunc)
+    return rulefunc
+
+
+@rule
+def jpeg(h):
     """JPEG data with JFIF or Exif markers; and raw JPEG"""
     if h[6:10] in (b'JFIF', b'Exif'):
         return 'jpeg'
     elif h[:4] == b'\xff\xd8\xff\xdb':
         return 'jpeg'
 
-tests.append(_test_jpeg)
 
-def _test_png(h):
+@rule
+def png(h):
     if h.startswith(b'\211PNG\r\n\032\n'):
         return 'png'
 
-tests.append(_test_png)
 
-def _test_gif(h):
+@rule
+def gif(h):
     """GIF ('87 and '89 variants)"""
     if h[:6] in (b'GIF87a', b'GIF89a'):
         return 'gif'
 
-tests.append(_test_gif)
 
-def _test_tiff(h):
+@rule
+def tiff(h):
     """TIFF (can be in Motorola or Intel byte order)"""
     if h[:2] in (b'MM', b'II'):
         return 'tiff'
 
-tests.append(_test_tiff)
 
-def _test_rgb(h):
+@rule
+def rgb(h):
     """SGI image library"""
     if h.startswith(b'\001\332'):
         return 'rgb'
 
-tests.append(_test_rgb)
 
-def _test_pbm(h):
+@rule
+def pbm(h):
     """PBM (portable bitmap)"""
     if len(h) >= 3 and \
-        h[0] == ord(b'P') and h[1] in b'14' and h[2] in b' \t\n\r':
+            h[0] == ord(b'P') and h[1] in b'14' and h[2] in b' \t\n\r':
         return 'pbm'
 
-tests.append(_test_pbm)
 
-def _test_pgm(h):
+@rule
+def pgm(h):
     """PGM (portable graymap)"""
     if len(h) >= 3 and \
-        h[0] == ord(b'P') and h[1] in b'25' and h[2] in b' \t\n\r':
+            h[0] == ord(b'P') and h[1] in b'25' and h[2] in b' \t\n\r':
         return 'pgm'
 
-tests.append(_test_pgm)
 
-def _test_ppm(h):
+@rule
+def ppm(h):
     """PPM (portable pixmap)"""
     if len(h) >= 3 and \
-        h[0] == ord(b'P') and h[1] in b'36' and h[2] in b' \t\n\r':
+            h[0] == ord(b'P') and h[1] in b'36' and h[2] in b' \t\n\r':
         return 'ppm'
 
-tests.append(_test_ppm)
 
-def _test_rast(h):
+@rule
+def rast(h):
     """Sun raster file"""
     if h.startswith(b'\x59\xA6\x6A\x95'):
         return 'rast'
 
-tests.append(_test_rast)
 
+@rule
 def _test_xbm(h):
     """X bitmap (X10 or X11)"""
     if h.startswith(b'#define '):
         return 'xbm'
 
-tests.append(_test_xbm)
 
-def _test_bmp(h):
+@rule
+def bmp(h):
     if h.startswith(b'BM'):
         return 'bmp'
 
-tests.append(_test_bmp)
 
-def _test_webp(h):
+@rule
+def webp(h):
     if h.startswith(b'RIFF') and h[8:12] == b'WEBP':
         return 'webp'
 
-tests.append(_test_webp)
 
+@rule
 def _test_exr(h):
     if h.startswith(b'\x76\x2f\x31\x01'):
         return 'exr'
-
-tests.append(_test_exr)
-
-
-class MIMEImage(MIMENonMultipart):
-    """Class for generating image/* type MIME documents."""
-
-    def __init__(self, _imagedata, _subtype=None,
-                 _encoder=encoders.encode_base64, *, policy=None, **_params):
-        """Create an image/* type MIME document.
-
-        _imagedata is a string containing the raw image data.  If the data
-        type can be detected (jpeg, png, gif, tiff, rgb, pbm, pgm, ppm,
-        rast, xbm, bmp, webp, and exr attempted), then the subtype will be
-        automatically included in the Content-Type header. Otherwise, you can
-        specify the specific image subtype via the _subtype parameter.
-
-        _encoder is a function which will perform the actual encoding for
-        transport of the image data.  It takes one argument, which is this
-        Image instance.  It should use get_payload() and set_payload() to
-        change the payload to the encoded form.  It should also add any
-        Content-Transfer-Encoding or other headers to the message as
-        necessary.  The default encoding is Base64.
-
-        Any additional keyword arguments are passed to the base class
-        constructor, which turns them into parameters on the Content-Type
-        header.
-        """
-        if _subtype is None:
-            if (_subtype := _what(_imagedata)) is None:
-                raise TypeError('Could not guess image MIME subtype')
-        MIMENonMultipart.__init__(self, 'image', _subtype, policy=policy,
-                                  **_params)
-        self.set_payload(_imagedata)
-        _encoder(self)
@@ -0,0 +1,3 @@
+P4
+16 16
+�������[�a_�X������?��
@@ -0,0 +1,6 @@
+#define python_width 16
+#define python_height 16
+static char python_bits[] = {
+  0xDF, 0xFE, 0x8F, 0xFD, 0x5F, 0xFB, 0xAB, 0xFE, 0xB5, 0x8D, 0xDA, 0x8F, 
+  0xA5, 0x86, 0xFA, 0x83, 0x1A, 0x80, 0x0D, 0x80, 0x0D, 0x80, 0x0F, 0xE0, 
+  0x0F, 0xF8, 0x0F, 0xF8, 0x0F, 0xFC, 0xFF, 0xFF, };
@@ -798,7 +798,7 @@ def test_unicode_body_defaults_to_utf8_encoding(self):
 class TestEncoders(unittest.TestCase):
 
     def test_EncodersEncode_base64(self):
-        with openfile('PyBanner048.gif', 'rb') as fp:
+        with openfile('python.gif', 'rb') as fp:
             bindata = fp.read()
         mimed = email.mime.image.MIMEImage(bindata)
         base64ed = mimed.get_payload()
@@ -1555,24 +1555,44 @@ def test_add_header(self):
 
 # Test the basic MIMEImage class
 class TestMIMEImage(unittest.TestCase):
-    def setUp(self):
-        with openfile('PyBanner048.gif', 'rb') as fp:
+    def _make_image(self, ext):
+        with openfile(f'python.{ext}', 'rb') as fp:
             self._imgdata = fp.read()
         self._im = MIMEImage(self._imgdata)
 
     def test_guess_minor_type(self):
-        self.assertEqual(self._im.get_content_type(), 'image/gif')
+        for ext, subtype in {
+            'bmp': None,
+            'exr': None,
+            'gif': None,
+            'jpg': 'jpeg',
+            'pbm': None,
+            'pgm': None,
+            'png': None,
+            'ppm': None,
+            'ras': 'rast',
+            'sgi': 'rgb',
+            'tiff': None,
+            'webp': None,
+            'xbm': None,
+        }.items():
+            self._make_image(ext)
+            subtype = ext if subtype is None else subtype
+            self.assertEqual(self._im.get_content_type(), f'image/{subtype}')
 
     def test_encoding(self):
+        self._make_image('gif')
         payload = self._im.get_payload()
         self.assertEqual(base64.decodebytes(bytes(payload, 'ascii')),
-                self._imgdata)
+                         self._imgdata)
 
     def test_checkSetMinor(self):
+        self._make_image('gif')
         im = MIMEImage(self._imgdata, 'fish')
         self.assertEqual(im.get_content_type(), 'image/fish')
 
     def test_add_header(self):
+        self._make_image('gif')
         eq = self.assertEqual
         self._im.add_header('Content-Disposition', 'attachment',
                             filename='dingusfish.gif')
@@ -1747,7 +1767,7 @@ def test_utf8_input_no_charset(self):
 # Test complicated multipart/* messages
 class TestMultipart(TestEmailBase):
     def setUp(self):
-        with openfile('PyBanner048.gif', 'rb') as fp:
+        with openfile('python.gif', 'rb') as fp:
             data = fp.read()
         container = MIMEBase('multipart', 'mixed', boundary='BOUNDARY')
         image = MIMEImage(data, name='dingusfish.gif')
@@ -3444,7 +3464,7 @@ def test_BytesGenerator_linend_with_non_ascii(self):
     def test_mime_classes_policy_argument(self):
         with openfile('audiotest.au', 'rb') as fp:
             audiodata = fp.read()
-        with openfile('PyBanner048.gif', 'rb') as fp:
+        with openfile('python.gif', 'rb') as fp:
             bindata = fp.read()
         classes = [
             (MIMEApplication, ('',)),