Added to_data_uri method for Image class. (#2227)

* Added to_data_uri method and path_to_data_uri classmethod for Image class.

* Removed path_to_data_uri classmethod and modified _get_mime_type to use mimetypes.guess_type function instead of hardcoded dictionary.

* Register image/webp with mimetypes before guess_type to support WEBP mimetype detection on Python 3.10.

* Improved branch coverage for Image.to_data_uri.

* Added Image._to_data_uri example in the docs.
This commit is contained in:
Harshith Thota 2025-11-01 21:10:08 +05:30 committed by GitHub
commit 321f404046
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 75 additions and 30 deletions

View file

@ -115,6 +115,7 @@ For small icons or when you want to embed the icon directly, use data URIs:
```python
from mcp.types import Icon
from fastmcp.utilities.types import Image
# SVG icon as data URI
svg_icon = Icon(
@ -126,4 +127,13 @@ svg_icon = Icon(
def my_tool() -> str:
"""A tool with an embedded SVG icon."""
return "result"
# Generating a data URI from a local image file.
img = Image(path="./assets/brand/favicon.png")
icon = Icon(src=img.to_data_uri())
@mcp.tool(icons=[icon])
def file_icon_tool() -> str:
"""A tool with an icon generated from a local file."""
return "result"
```

View file

@ -190,34 +190,33 @@ class Image:
if path is not None and data is not None:
raise ValueError("Only one of path or data can be provided")
self.path = Path(os.path.expandvars(str(path))).expanduser() if path else None
self.path = self._get_expanded_path(path)
self.data = data
self._format = format
self._mime_type = self._get_mime_type()
self.annotations = annotations
@staticmethod
def _get_expanded_path(path: str | Path | None) -> Path | None:
"""Expand environment variables and user home in path."""
return Path(os.path.expandvars(str(path))).expanduser() if path else None
def _get_mime_type(self) -> str:
"""Get MIME type from format or guess from file extension."""
if self._format:
return f"image/{self._format.lower()}"
if self.path:
suffix = self.path.suffix.lower()
return {
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
}.get(suffix, "application/octet-stream")
# Workaround for WEBP in Py3.10
mimetypes.add_type("image/webp", ".webp")
resp = mimetypes.guess_type(self.path, strict=False)
if resp and resp[0] is not None:
return resp[0]
return "application/octet-stream"
return "image/png" # default for raw binary data
def to_image_content(
self,
mime_type: str | None = None,
annotations: Annotations | None = None,
) -> mcp.types.ImageContent:
"""Convert to MCP ImageContent."""
def _get_data(self) -> str:
"""Get raw image data as base64-encoded string."""
if self.path:
with open(self.path, "rb") as f:
data = base64.b64encode(f.read()).decode()
@ -225,6 +224,15 @@ class Image:
data = base64.b64encode(self.data).decode()
else:
raise ValueError("No image data available")
return data
def to_image_content(
self,
mime_type: str | None = None,
annotations: Annotations | None = None,
) -> mcp.types.ImageContent:
"""Convert to MCP ImageContent."""
data = self._get_data()
return mcp.types.ImageContent(
type="image",
@ -233,6 +241,11 @@ class Image:
annotations=annotations or self.annotations,
)
def to_data_uri(self, mime_type: str | None = None) -> str:
"""Get image as a data URI."""
data = self._get_data()
return f"data:{mime_type or self._mime_type};base64,{data}"
class Audio:
"""Helper class for returning audio from tools."""

View file

@ -177,22 +177,23 @@ class TestImage:
):
Image(path="test.png", data=b"test")
def test_get_mime_type_from_path(self, tmp_path):
@pytest.mark.parametrize(
"extension,mime_type",
[
(".png", "image/png"),
(".jpg", "image/jpeg"),
(".jpeg", "image/jpeg"),
(".gif", "image/gif"),
(".webp", "image/webp"),
(".unknown", "application/octet-stream"),
],
)
def test_get_mime_type_from_path(self, tmp_path, extension, mime_type):
"""Test MIME type detection from file extension."""
extensions = {
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
".unknown": "application/octet-stream",
}
for ext, mime in extensions.items():
path = tmp_path / f"test{ext}"
path.write_bytes(b"fake image data")
img = Image(path=path)
assert img._mime_type == mime
path = tmp_path / f"test{extension}"
path.write_bytes(b"fake image data")
img = Image(path=path)
assert img._mime_type == mime_type
def test_to_image_content(self, tmp_path, monkeypatch):
"""Test conversion to ImageContent."""
@ -227,6 +228,27 @@ class TestImage:
with pytest.raises(ValueError, match="No image data available"):
img.to_image_content()
@pytest.mark.parametrize(
"mime_type,fname,expected_mime",
[
(None, "test.png", "image/png"),
("image/jpeg", "test.unknown", "image/jpeg"),
],
)
def test_to_data_uri(self, tmp_path, mime_type, fname, expected_mime):
"""Test conversion to data URI."""
img_path = tmp_path / fname
test_data = b"fake image data"
img_path.write_bytes(test_data)
img = Image(path=img_path)
data_uri = img.to_data_uri(mime_type=mime_type)
expected_data_uri = (
f"data:{expected_mime};base64,{base64.b64encode(test_data).decode()}"
)
assert data_uri == expected_data_uri
class TestAudio:
def test_audio_initialization_with_path(self):