Switched to auto parsing of mime_type

m-messer · m-messer · commit cdf120cc8687 · 2026-01-28T13:37:59.000Z
diff --git a/lf_toolkit/evaluation/image_upload.py b/lf_toolkit/evaluation/image_upload.py
@@ -13,22 +13,16 @@
     'image/png': ['PNG'],
     'image/gif': ['GIF'],
     'image/bmp': ['BMP'],
-    'image/webp': ['WEBP'],
-    'image/tiff': ['TIFF', 'TIF'],
-    'image/x-icon': ['ICO'],
 }
 
-FORMAT_TO_EXTENSION: Dict[str, List[str]] = {
-    'JPEG': ['.jpg', '.jpeg', '.jpe'],
-    'PNG': ['.png'],
-    'GIF': ['.gif'],
-    'BMP': ['.bmp'],
-    'WEBP': ['.webp'],
-    'TIFF': ['.tiff', '.tif'],
-    'ICO': ['.ico'],
+FORMAT_TO_MIME: Dict[str, str] = {
+    'JPEG': 'image/jpeg',
+    'JPG': 'image/jpeg',
+    'PNG': 'image/png',
+    'GIF': 'image/gif',
+    "bmp": 'image/bmp'
 }
 
-
 class ImageUploadError(Exception):
     """Custom exception for image upload failures"""
     pass
@@ -57,50 +51,6 @@ def generate_file_name(img: Image.Image) -> str:
     format_ext: str = img.format.lower() if img.format else 'png'
     return f"{unique_id}.{format_ext}"
 
-
-def validate_mime_type(mime_type: str, img: Image.Image, filename: str) -> bool:
-    """Validate MIME type against image format and filename
-
-    Args:
-        mime_type: MIME type string to validate
-        img: PIL Image object
-        filename: Name of the file
-
-    Returns:
-        True if validation passes
-
-    Raises:
-        InvalidMimeTypeError: If MIME type is invalid or doesn't match image
-    """
-    if mime_type not in MIME_TO_FORMAT:
-        raise InvalidMimeTypeError(
-            f"Invalid MIME type '{mime_type}'. "
-            f"Supported types: {', '.join(MIME_TO_FORMAT.keys())}"
-        )
-
-    img_format: Optional[str] = img.format.upper() if img.format else None
-
-    if img_format:
-        allowed_formats: List[str] = MIME_TO_FORMAT[mime_type]
-        if img_format not in allowed_formats:
-            raise InvalidMimeTypeError(
-                f"MIME type '{mime_type}' does not match image format '{img_format}'. "
-                f"Expected formats for {mime_type}: {', '.join(allowed_formats)}"
-            )
-
-    file_ext: str = filename[filename.rfind('.'):].lower()
-
-    if img_format and img_format in FORMAT_TO_EXTENSION:
-        valid_extensions: List[str] = FORMAT_TO_EXTENSION[img_format]
-        if file_ext not in valid_extensions:
-            raise InvalidMimeTypeError(
-                f"File extension '{file_ext}' does not match format '{img_format}'. "
-                f"Expected extensions: {', '.join(valid_extensions)}"
-            )
-
-    return True
-
-
 def get_s3_bucket_uri() -> str:
     """Get S3 bucket URI from environment variable"""
     s3_uri: Optional[str] = os.getenv('S3_BUCKET_URI')
@@ -113,12 +63,11 @@ def get_s3_bucket_uri() -> str:
     return s3_uri
 
 
-def upload_image(img: Image.Image, mime_type: str) -> str:
+def upload_image(img: Image.Image) -> str:
     """Upload PIL image with comprehensive MIME type validation
 
     Args:
         img: PIL Image object to upload
-        mime_type: MIME type for the upload
 
     Returns:
         JSON response from the server as a dictionary
@@ -134,10 +83,13 @@ def upload_image(img: Image.Image, mime_type: str) -> str:
 
         filename: str = generate_file_name(img)
 
-        validate_mime_type(mime_type, img, filename)
-
         full_url = base_url + filename
 
+        if img.format is None:
+            img.format = 'PNG'
+
+        mime_type = FORMAT_TO_MIME[img.format.upper()]
+
         buffer: BytesIO = BytesIO()
         img_format: str = img.format if img.format else 'PNG'
         img.save(buffer, format=img_format)
@@ -169,5 +121,5 @@ def upload_image(img: Image.Image, mime_type: str) -> str:
     img.format = 'JPEG'
 
     # Execute
-    result = upload_image(img, 'image/jpeg')
+    result = upload_image(img)
     print(result)
diff --git a/tests/evaluation/image_upload_test.py b/tests/evaluation/image_upload_test.py
@@ -7,14 +7,12 @@
 # Import the module to test
 from lf_toolkit.evaluation.image_upload import (
     generate_file_name,
-    validate_mime_type,
     get_s3_bucket_uri,
     upload_image,
     ImageUploadError,
     InvalidMimeTypeError,
     MissingEnvironmentVariableError,
     MIME_TO_FORMAT,
-    FORMAT_TO_EXTENSION
 )
 
 
@@ -65,89 +63,6 @@ def test_generate_file_name_unique(self):
         assert filename1 != filename2
 
 
-class TestValidateMimeType:
-    """Test suite for validate_mime_type function"""
-
-    def test_valid_jpeg_mime_type(self):
-        """Test validation with valid JPEG MIME type"""
-        img = Mock(spec=Image.Image)
-        img.format = 'JPEG'
-
-        result = validate_mime_type('image/jpeg', img, 'test.jpg')
-        assert result is True
-
-    def test_valid_png_mime_type(self):
-        """Test validation with valid PNG MIME type"""
-        img = Mock(spec=Image.Image)
-        img.format = 'PNG'
-
-        result = validate_mime_type('image/png', img, 'test.png')
-        assert result is True
-
-    def test_invalid_mime_type(self):
-        """Test validation with unsupported MIME type"""
-        img = Mock(spec=Image.Image)
-        img.format = 'PNG'
-
-        with pytest.raises(InvalidMimeTypeError) as exc_info:
-            validate_mime_type('image/invalid', img, 'test.png')
-
-        assert "Invalid MIME type 'image/invalid'" in str(exc_info.value)
-
-    def test_mime_type_format_mismatch(self):
-        """Test validation when MIME type doesn't match image format"""
-        img = Mock(spec=Image.Image)
-        img.format = 'PNG'
-
-        with pytest.raises(InvalidMimeTypeError) as exc_info:
-            validate_mime_type('image/jpeg', img, 'test.png')
-
-        assert "does not match image format 'PNG'" in str(exc_info.value)
-
-    def test_extension_format_mismatch(self):
-        """Test validation when file extension doesn't match format"""
-        img = Mock(spec=Image.Image)
-        img.format = 'JPEG'
-
-        with pytest.raises(InvalidMimeTypeError) as exc_info:
-            validate_mime_type('image/jpeg', img, 'test.png')
-
-        assert "File extension '.png' does not match format 'JPEG'" in str(exc_info.value)
-
-    def test_valid_with_no_image_format(self):
-        """Test validation when image has no format attribute"""
-        img = Mock(spec=Image.Image)
-        img.format = None
-
-        # Should not raise when format is None
-        result = validate_mime_type('image/png', img, 'test.png')
-        assert result is True
-
-    def test_valid_webp_mime_type(self):
-        """Test validation with valid WEBP MIME type"""
-        img = Mock(spec=Image.Image)
-        img.format = 'WEBP'
-
-        result = validate_mime_type('image/webp', img, 'test.webp')
-        assert result is True
-
-    def test_jpeg_with_jpg_extension(self):
-        """Test JPEG image with .jpg extension"""
-        img = Mock(spec=Image.Image)
-        img.format = 'JPEG'
-
-        result = validate_mime_type('image/jpeg', img, 'photo.jpg')
-        assert result is True
-
-    def test_jpeg_with_jpeg_extension(self):
-        """Test JPEG image with .jpeg extension"""
-        img = Mock(spec=Image.Image)
-        img.format = 'JPEG'
-
-        result = validate_mime_type('image/jpeg', img, 'photo.jpeg')
-        assert result is True
-
-
 class TestGetS3BucketUri:
     """Test suite for get_s3_bucket_uri function"""
 
@@ -201,7 +116,7 @@ def test_successful_upload(self, mock_uuid, mock_getenv, mock_put):
         img.format = 'JPEG'
 
         # Execute
-        result = upload_image(img, 'image/jpeg')
+        result = upload_image(img)
 
         # Verify response
         assert result == 'https://s3.amazonaws.com/my-bucket/12345678-1234-5678-1234-567812345678.jpeg'
@@ -223,7 +138,7 @@ def test_upload_with_png_image(self, mock_uuid, mock_getenv, mock_put):
         img = Image.new('RGBA', (50, 50), color=(0, 255, 0, 128))
         img.format = 'PNG'
 
-        result = upload_image(img, 'image/png')
+        result = upload_image(img)
 
         assert result == 'https://storage.example.com/aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee.png'
 
@@ -236,18 +151,7 @@ def test_upload_missing_s3_uri(self, mock_getenv):
         img.format = 'JPEG'
 
         with pytest.raises(MissingEnvironmentVariableError):
-            upload_image(img, 'image/jpeg')
-
-    @patch('lf_toolkit.evaluation.image_upload.os.getenv')
-    def test_upload_invalid_mime_type(self, mock_getenv):
-        """Test upload fails with invalid MIME type"""
-        mock_getenv.return_value = 'https://s3.amazonaws.com/bucket'
-
-        img = Image.new('RGB', (100, 100))
-        img.format = 'JPEG'
-
-        with pytest.raises(InvalidMimeTypeError):
-            upload_image(img, 'image/invalid')
+            upload_image(img)
 
     @patch('lf_toolkit.evaluation.image_upload.requests.put')
     @patch('lf_toolkit.evaluation.image_upload.os.getenv')
@@ -266,7 +170,7 @@ def test_upload_server_error(self, mock_uuid, mock_getenv, mock_put):
         img.format = 'JPEG'
 
         with pytest.raises(ImageUploadError) as exc_info:
-            upload_image(img, 'image/jpeg')
+            upload_image(img)
 
         assert "Upload failed with status code 500" in str(exc_info.value)
 
@@ -284,7 +188,7 @@ def test_upload_network_error(self, mock_uuid, mock_getenv, mock_put):
         img.format = 'JPEG'
 
         with pytest.raises(ImageUploadError) as exc_info:
-            upload_image(img, 'image/jpeg')
+            upload_image(img)
 
         assert "Network error" in str(exc_info.value)
 
@@ -302,24 +206,10 @@ def test_upload_timeout_error(self, mock_uuid, mock_getenv, mock_put):
         img.format = 'JPEG'
 
         with pytest.raises(ImageUploadError) as exc_info:
-            upload_image(img, 'image/jpeg')
+            upload_image(img)
 
         assert "Network error" in str(exc_info.value)
 
-    @patch('lf_toolkit.evaluation.image_upload.requests.put')
-    @patch('lf_toolkit.evaluation.image_upload.os.getenv')
-    @patch('lf_toolkit.evaluation.image_upload.uuid.uuid4')
-    def test_upload_mime_type_mismatch(self, mock_uuid, mock_getenv, mock_put):
-        """Test upload fails when MIME type doesn't match image format"""
-        mock_uuid.return_value = uuid.UUID('12345678-1234-5678-1234-567812345678')
-        mock_getenv.return_value = 'https://s3.amazonaws.com/bucket'
-
-        img = Image.new('RGB', (100, 100))
-        img.format = 'PNG'
-
-        with pytest.raises(InvalidMimeTypeError):
-            upload_image(img, 'image/jpeg')
-
     @patch('lf_toolkit.evaluation.image_upload.requests.put')
     @patch('lf_toolkit.evaluation.image_upload.os.getenv')
     @patch('lf_toolkit.evaluation.image_upload.uuid.uuid4')
@@ -335,7 +225,7 @@ def test_upload_image_no_format(self, mock_uuid, mock_getenv, mock_put):
         img = Image.new('RGB', (100, 100))
         img.format = None
 
-        result = upload_image(img, 'image/png')
+        result = upload_image(img)
 
         assert result == 'https://s3.amazonaws.com/bucket/12345678-1234-5678-1234-567812345678.png'
 
@@ -354,7 +244,7 @@ def test_upload_verifies_correct_file_uploaded(self, mock_uuid, mock_getenv, moc
         img = Image.new('RGB', (100, 100), color='blue')
         img.format = 'JPEG'
 
-        upload_image(img, 'image/jpeg')
+        upload_image(img)
 
 
 class TestExceptionHierarchy:
@@ -385,28 +275,5 @@ def test_invalid_mime_type_error_caught_as_image_upload_error(self):
             raise InvalidMimeTypeError("Invalid MIME")
 
 
-class TestConstants:
-    """Test suite for module constants"""
-
-    def test_mime_to_format_has_expected_types(self):
-        """Test that MIME_TO_FORMAT contains expected image types"""
-        assert 'image/jpeg' in MIME_TO_FORMAT
-        assert 'image/png' in MIME_TO_FORMAT
-        assert 'image/gif' in MIME_TO_FORMAT
-        assert 'image/webp' in MIME_TO_FORMAT
-
-    def test_format_to_extension_has_expected_formats(self):
-        """Test that FORMAT_TO_EXTENSION contains expected formats"""
-        assert 'JPEG' in FORMAT_TO_EXTENSION
-        assert 'PNG' in FORMAT_TO_EXTENSION
-        assert 'GIF' in FORMAT_TO_EXTENSION
-        assert 'WEBP' in FORMAT_TO_EXTENSION
-
-    def test_jpeg_has_multiple_extensions(self):
-        """Test that JPEG format has multiple valid extensions"""
-        assert '.jpg' in FORMAT_TO_EXTENSION['JPEG']
-        assert '.jpeg' in FORMAT_TO_EXTENSION['JPEG']
-
-
 if __name__ == '__main__':
     pytest.main([__file__, '-v'])
