improve messages coverage (#1704)

Author: samuelcolvin

pydantic_ai_slim/pydantic_ai/messages.py

@@ -83,7 +83,7 @@ class VideoUrl:
     """Type identifier, this is available on all parts as a discriminator."""
 
     @property
-    def media_type(self) -> VideoMediaType:  # pragma: lax no cover
+    def media_type(self) -> VideoMediaType:
         """Return the media type of the video, based on the url."""
         if self.url.endswith('.mkv'):
             return 'video/x-matroska'
@@ -110,7 +110,7 @@ def format(self) -> VideoFormat:
 
         The choice of supported formats were based on the Bedrock Converse API. Other APIs don't require to use a format.
         """
-        return _video_format(self.media_type)
+        return _video_format_lookup[self.media_type]
 
 
 @dataclass
@@ -133,6 +133,11 @@ def media_type(self) -> AudioMediaType:
         else:
             raise ValueError(f'Unknown audio file extension: {self.url}')
 
+    @property
+    def format(self) -> AudioFormat:
+        """The file format of the audio file."""
+        return _audio_format_lookup[self.media_type]
+
 
 @dataclass
 class ImageUrl:
@@ -164,7 +169,7 @@ def format(self) -> ImageFormat:
 
         The choice of supported formats were based on the Bedrock Converse API. Other APIs don't require to use a format.
         """
-        return _image_format(self.media_type)
+        return _image_format_lookup[self.media_type]
 
 
 @dataclass
@@ -182,7 +187,7 @@ def media_type(self) -> str:
         """Return the media type of the document, based on the url."""
         type_, _ = guess_type(self.url)
         if type_ is None:
-            raise RuntimeError(f'Unknown document file extension: {self.url}')
+            raise ValueError(f'Unknown document file extension: {self.url}')
         return type_
 
     @property
@@ -191,7 +196,11 @@ def format(self) -> DocumentFormat:
 
         The choice of supported formats were based on the Bedrock Converse API. Other APIs don't require to use a format.
         """
-        return _document_format(self.media_type)
+        media_type = self.media_type
+        try:
+            return _document_format_lookup[media_type]
+        except KeyError as e:
+            raise ValueError(f'Unknown document media type: {media_type}') from e
 
 
 @dataclass
@@ -225,93 +234,58 @@ def is_video(self) -> bool:
     @property
     def is_document(self) -> bool:
         """Return `True` if the media type is a document type."""
-        return self.media_type in {
-            'application/pdf',
-            'text/plain',
-            'text/csv',
-            'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
-            'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
-            'text/html',
-            'text/markdown',
-            'application/vnd.ms-excel',
-        }
+        return self.media_type in _document_format_lookup
 
     @property
     def format(self) -> str:
         """The file format of the binary content."""
-        if self.is_audio:
-            if self.media_type == 'audio/mpeg':
-                return 'mp3'
-            elif self.media_type == 'audio/wav':
-                return 'wav'
-        elif self.is_image:
-            return _image_format(self.media_type)
-        elif self.is_document:
-            return _document_format(self.media_type)
-        elif self.is_video:
-            return _video_format(self.media_type)
-        raise ValueError(f'Unknown media type: {self.media_type}')
+        try:
+            if self.is_audio:
+                return _audio_format_lookup[self.media_type]
+            elif self.is_image:
+                return _image_format_lookup[self.media_type]
+            elif self.is_video:
+                return _video_format_lookup[self.media_type]
+            else:
+                return _document_format_lookup[self.media_type]
+        except KeyError as e:
+            raise ValueError(f'Unknown media type: {self.media_type}') from e
 
 
 UserContent: TypeAlias = 'str | ImageUrl | AudioUrl | DocumentUrl | VideoUrl | BinaryContent'
 
 # Ideally this would be a Union of types, but Python 3.9 requires it to be a string, and strings don't work with `isinstance``.
 MultiModalContentTypes = (ImageUrl, AudioUrl, DocumentUrl, VideoUrl, BinaryContent)
-
-
-def _document_format(media_type: str) -> DocumentFormat:
-    if media_type == 'application/pdf':
-        return 'pdf'
-    elif media_type == 'text/plain':
-        return 'txt'
-    elif media_type == 'text/csv':
-        return 'csv'
-    elif media_type == 'application/vnd.openxmlformats-officedocument.wordprocessingml.document':
-        return 'docx'
-    elif media_type == 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet':
-        return 'xlsx'
-    elif media_type == 'text/html':
-        return 'html'
-    elif media_type == 'text/markdown':
-        return 'md'
-    elif media_type == 'application/vnd.ms-excel':
-        return 'xls'
-    else:
-        raise ValueError(f'Unknown document media type: {media_type}')
-
-
-def _image_format(media_type: str) -> ImageFormat:
-    if media_type == 'image/jpeg':
-        return 'jpeg'
-    elif media_type == 'image/png':
-        return 'png'
-    elif media_type == 'image/gif':
-        return 'gif'
-    elif media_type == 'image/webp':
-        return 'webp'
-    else:
-        raise ValueError(f'Unknown image media type: {media_type}')
-
-
-def _video_format(media_type: str) -> VideoFormat:
-    if media_type == 'video/x-matroska':
-        return 'mkv'
-    elif media_type == 'video/quicktime':
-        return 'mov'
-    elif media_type == 'video/mp4':
-        return 'mp4'
-    elif media_type == 'video/webm':
-        return 'webm'
-    elif media_type == 'video/x-flv':
-        return 'flv'
-    elif media_type == 'video/mpeg':
-        return 'mpeg'
-    elif media_type == 'video/x-ms-wmv':
-        return 'wmv'
-    elif media_type == 'video/3gpp':
-        return 'three_gp'
-    else:  # pragma: no cover
-        raise ValueError(f'Unknown video media type: {media_type}')
+_document_format_lookup: dict[str, DocumentFormat] = {
+    'application/pdf': 'pdf',
+    'text/plain': 'txt',
+    'text/csv': 'csv',
+    'application/vnd.openxmlformats-officedocument.wordprocessingml.document': 'docx',
+    'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet': 'xlsx',
+    'text/html': 'html',
+    'text/markdown': 'md',
+    'application/vnd.ms-excel': 'xls',
+}
+_audio_format_lookup: dict[str, AudioFormat] = {
+    'audio/mpeg': 'mp3',
+    'audio/wav': 'wav',
+}
+_image_format_lookup: dict[str, ImageFormat] = {
+    'image/jpeg': 'jpeg',
+    'image/png': 'png',
+    'image/gif': 'gif',
+    'image/webp': 'webp',
+}
+_video_format_lookup: dict[str, VideoFormat] = {
+    'video/x-matroska': 'mkv',
+    'video/quicktime': 'mov',
+    'video/mp4': 'mp4',
+    'video/webm': 'webm',
+    'video/x-flv': 'flv',
+    'video/mpeg': 'mpeg',
+    'video/x-ms-wmv': 'wmv',
+    'video/3gpp': 'three_gp',
+}
 
 
 @dataclass

tests/test_messages.py

@@ -1,6 +1,6 @@
 import pytest
 
-from pydantic_ai.messages import BinaryContent, DocumentUrl, ImageUrl, VideoUrl
+from pydantic_ai.messages import AudioUrl, BinaryContent, DocumentUrl, ImageUrl, VideoUrl
 
 
 def test_image_url():
@@ -20,7 +20,7 @@ def test_video_url():
 
 
 def test_document_url():
-    with pytest.raises(RuntimeError, match='Unknown document file extension: https://example.com/document.potato'):
+    with pytest.raises(ValueError, match='Unknown document file extension: https://example.com/document.potato'):
         document_url = DocumentUrl(url='https://example.com/document.potato')
         document_url.media_type
 
@@ -93,3 +93,163 @@ def test_binary_content_document(media_type: str, format: str):
     binary_content = BinaryContent(data=b'Hello, world!', media_type=media_type)
     assert binary_content.is_document
     assert binary_content.format == format
+
+
+@pytest.mark.parametrize(
+    'audio_url,media_type,format',
+    [
+        pytest.param(AudioUrl('foobar.mp3'), 'audio/mpeg', 'mp3', id='mp3'),
+        pytest.param(AudioUrl('foobar.wav'), 'audio/wav', 'wav', id='wav'),
+    ],
+)
+def test_audio_url(audio_url: AudioUrl, media_type: str, format: str):
+    assert audio_url.media_type == media_type
+    assert audio_url.format == format
+
+
+def test_audio_url_invalid():
+    with pytest.raises(ValueError, match='Unknown audio file extension: foobar.potato'):
+        AudioUrl('foobar.potato').media_type
+
+
+@pytest.mark.parametrize(
+    'image_url,media_type,format',
+    [
+        pytest.param(ImageUrl('foobar.jpg'), 'image/jpeg', 'jpeg', id='jpg'),
+        pytest.param(ImageUrl('foobar.jpeg'), 'image/jpeg', 'jpeg', id='jpeg'),
+        pytest.param(ImageUrl('foobar.png'), 'image/png', 'png', id='png'),
+        pytest.param(ImageUrl('foobar.gif'), 'image/gif', 'gif', id='gif'),
+        pytest.param(ImageUrl('foobar.webp'), 'image/webp', 'webp', id='webp'),
+    ],
+)
+def test_image_url_formats(image_url: ImageUrl, media_type: str, format: str):
+    assert image_url.media_type == media_type
+    assert image_url.format == format
+
+
+def test_image_url_invalid():
+    with pytest.raises(ValueError, match='Unknown image file extension: foobar.potato'):
+        ImageUrl('foobar.potato').media_type
+
+    with pytest.raises(ValueError, match='Unknown image file extension: foobar.potato'):
+        ImageUrl('foobar.potato').format
+
+
+@pytest.mark.parametrize(
+    'document_url,media_type,format',
+    [
+        pytest.param(DocumentUrl('foobar.pdf'), 'application/pdf', 'pdf', id='pdf'),
+        pytest.param(DocumentUrl('foobar.txt'), 'text/plain', 'txt', id='txt'),
+        pytest.param(DocumentUrl('foobar.csv'), 'text/csv', 'csv', id='csv'),
+        pytest.param(
+            DocumentUrl('foobar.docx'),
+            'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+            'docx',
+            id='docx',
+        ),
+        pytest.param(
+            DocumentUrl('foobar.xlsx'),
+            'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
+            'xlsx',
+            id='xlsx',
+        ),
+        pytest.param(DocumentUrl('foobar.html'), 'text/html', 'html', id='html'),
+        pytest.param(DocumentUrl('foobar.md'), 'text/markdown', 'md', id='md'),
+        pytest.param(DocumentUrl('foobar.xls'), 'application/vnd.ms-excel', 'xls', id='xls'),
+    ],
+)
+def test_document_url_formats(document_url: DocumentUrl, media_type: str, format: str):
+    assert document_url.media_type == media_type
+    assert document_url.format == format
+
+
+def test_document_url_invalid():
+    with pytest.raises(ValueError, match='Unknown document file extension: foobar.potato'):
+        DocumentUrl('foobar.potato').media_type
+
+    with pytest.raises(ValueError, match='Unknown document media type: text/x-python'):
+        DocumentUrl('foobar.py').format
+
+
+def test_binary_content_unknown_media_type():
+    with pytest.raises(ValueError, match='Unknown media type: application/custom'):
+        binary_content = BinaryContent(data=b'Hello, world!', media_type='application/custom')
+        binary_content.format
+
+
+def test_binary_content_is_methods():
+    # Test that is_X returns False for non-matching media types
+    audio_content = BinaryContent(data=b'Hello, world!', media_type='audio/wav')
+    assert audio_content.is_audio is True
+    assert audio_content.is_image is False
+    assert audio_content.is_video is False
+    assert audio_content.is_document is False
+    assert audio_content.format == 'wav'
+
+    audio_content = BinaryContent(data=b'Hello, world!', media_type='audio/wrong')
+    assert audio_content.is_audio is True
+    assert audio_content.is_image is False
+    assert audio_content.is_video is False
+    assert audio_content.is_document is False
+    with pytest.raises(ValueError, match='Unknown media type: audio/wrong'):
+        audio_content.format
+
+    audio_content = BinaryContent(data=b'Hello, world!', media_type='image/wrong')
+    assert audio_content.is_audio is False
+    assert audio_content.is_image is True
+    assert audio_content.is_video is False
+    assert audio_content.is_document is False
+    with pytest.raises(ValueError, match='Unknown media type: image/wrong'):
+        audio_content.format
+
+    image_content = BinaryContent(data=b'Hello, world!', media_type='image/jpeg')
+    assert image_content.is_audio is False
+    assert image_content.is_image is True
+    assert image_content.is_video is False
+    assert image_content.is_document is False
+    assert image_content.format == 'jpeg'
+
+    video_content = BinaryContent(data=b'Hello, world!', media_type='video/mp4')
+    assert video_content.is_audio is False
+    assert video_content.is_image is False
+    assert video_content.is_video is True
+    assert video_content.is_document is False
+    assert video_content.format == 'mp4'
+
+    video_content = BinaryContent(data=b'Hello, world!', media_type='video/wrong')
+    assert video_content.is_audio is False
+    assert video_content.is_image is False
+    assert video_content.is_video is True
+    assert video_content.is_document is False
+    with pytest.raises(ValueError, match='Unknown media type: video/wrong'):
+        video_content.format
+
+    document_content = BinaryContent(data=b'Hello, world!', media_type='application/pdf')
+    assert document_content.is_audio is False
+    assert document_content.is_image is False
+    assert document_content.is_video is False
+    assert document_content.is_document is True
+    assert document_content.format == 'pdf'
+
+
+@pytest.mark.parametrize(
+    'video_url,media_type,format',
+    [
+        pytest.param(VideoUrl('foobar.mp4'), 'video/mp4', 'mp4', id='mp4'),
+        pytest.param(VideoUrl('foobar.mov'), 'video/quicktime', 'mov', id='mov'),
+        pytest.param(VideoUrl('foobar.mkv'), 'video/x-matroska', 'mkv', id='mkv'),
+        pytest.param(VideoUrl('foobar.webm'), 'video/webm', 'webm', id='webm'),
+        pytest.param(VideoUrl('foobar.flv'), 'video/x-flv', 'flv', id='flv'),
+        pytest.param(VideoUrl('foobar.mpeg'), 'video/mpeg', 'mpeg', id='mpeg'),
+        pytest.param(VideoUrl('foobar.wmv'), 'video/x-ms-wmv', 'wmv', id='wmv'),
+        pytest.param(VideoUrl('foobar.three_gp'), 'video/3gpp', 'three_gp', id='three_gp'),
+    ],
+)
+def test_video_url_formats(video_url: VideoUrl, media_type: str, format: str):
+    assert video_url.media_type == media_type
+    assert video_url.format == format
+
+
+def test_video_url_invalid():
+    with pytest.raises(ValueError, match='Unknown video file extension: foobar.potato'):
+        VideoUrl('foobar.potato').media_type