Add method for getting file extension

rkoopmans · rkoopmans · commit f175348e730f · 2022-09-20T13:48:25.000+02:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -8,6 +8,7 @@
     * Added python 3.11
 * Fixed tests on windows
 * Add methods for the transcoding and transformation API
+* Add a method for getting the file extension from a Result object
 
 ## 1.5.2
 Remove letsencrypt DST Root from ca bundle for openssl 1.0.0 compatibility
diff --git a/test/tinify_result_test.py b/test/tinify_result_test.py
@@ -35,6 +35,11 @@ def test_content_type_should_return_mime_type(self):
     def test_to_buffer_should_return_image_data(self):
         self.assertEqual(b'image data', self.result.to_buffer())
 
+    def test_extension(self):
+        self.assertEqual('png', self.result.extension)
+
+
+
 class TinifyResultWithoutMetaAndDataTest(TestHelper):
     def setUp(self):
         self.result = Result({}, None)
@@ -59,3 +64,6 @@ def test_content_type_should_return_none(self):
 
     def test_to_buffer_should_return_none(self):
         self.assertEqual(None, self.result.to_buffer())
+
+    def test_extension(self):
+        self.assertEqual(None, self.result.extension)
diff --git a/tinify/result.py b/tinify/result.py
@@ -27,6 +27,12 @@ def size(self):
     def media_type(self):
         return self._meta.get('Content-Type')
 
+    @property
+    def extension(self):
+        media_type = self._meta.get('Content-Type')
+        if media_type:
+            return media_type.split('/')[-1]
+
     @property
     def content_type(self):
         return self.media_type
