diff --git a/python/cog/__init__.py b/python/cog/__init__.py index 988af6e285..dd3bcedf4f 100644 --- a/python/cog/__init__.py +++ b/python/cog/__init__.py @@ -1,6 +1,9 @@ +import mimetypes + from pydantic import BaseModel from .base_predictor import BasePredictor +from .mimetypes_ext import install_mime_extensions from .server.scope import current_scope from .types import ( ConcatenateIterator, @@ -11,6 +14,8 @@ Secret, ) +install_mime_extensions(mimetypes) + try: from ._version import __version__ except ImportError: diff --git a/python/cog/mimetypes_ext.py b/python/cog/mimetypes_ext.py new file mode 100644 index 0000000000..d1a13f10cf --- /dev/null +++ b/python/cog/mimetypes_ext.py @@ -0,0 +1,21 @@ +import typing + +if typing.TYPE_CHECKING: + from typing_extensions import Protocol +else: + Protocol = object + + +class IMimeTypes(Protocol): + def add_type(self, type: str, ext: str, strict: bool = True) -> None: ... + + +def install_mime_extensions(mimetypes: IMimeTypes) -> None: + """ + Older versions of Python are missing the MIME types for more recent file formats + this function adds the missing MIME types to the mimetypes module. + """ + + # This could also be done by loading a mime.types file from disk using + # mimetypes.read_mime_types(). + mimetypes.add_type("image/webp", ".webp") diff --git a/python/tests/test_mimetypes_ext.py b/python/tests/test_mimetypes_ext.py new file mode 100644 index 0000000000..a1e5303321 --- /dev/null +++ b/python/tests/test_mimetypes_ext.py @@ -0,0 +1,18 @@ +from mimetypes import MimeTypes + +from cog.mimetypes_ext import install_mime_extensions + + +def test_webp_ext_support(): + # Assert on empty database. + mt = MimeTypes(filenames=tuple()) + assert mt.guess_type("image.webp") == (None, None) + install_mime_extensions(mt) + assert mt.guess_type("image.webp") == ("image/webp", None) + + # Assert global override + import mimetypes + + import cog # noqa: F401 + + assert mimetypes.guess_type("image.webp") == ("image/webp", None)