Skip to content

Commit

Permalink
Add support for image/webp to mimetypes package
Browse files Browse the repository at this point in the history
This has only recently been introduced in Python 3.13.0 and is currently
inconsistently implemented across different platforms. Confusingly webp
is supported in local development on macOS but not when building the
docker image of a cog model. This is either because it's not defined
in the system mime.types file of the Linux image or because a dev
dependency is manually adding it. I've not done the work to fully
understand which.

This commit introduces a function called in the init script for the cog
package that patches the global mimetypes registry to understand the
.webp extension and image/webp mime type. This will be a no-op on
systems that already understand the type.

This fixes a bug whereby files with the .webp extension are uploaded
to the `--upload-url` with the incorrect application/octet-stream header.
  • Loading branch information
aron committed Nov 19, 2024
1 parent 425d5a2 commit 5595562
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 0 deletions.
5 changes: 5 additions & 0 deletions python/cog/__init__.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -11,6 +14,8 @@
Secret,
)

install_mime_extensions(mimetypes)

try:
from ._version import __version__
except ImportError:
Expand Down
21 changes: 21 additions & 0 deletions python/cog/mimetypes_ext.py
Original file line number Diff line number Diff line change
@@ -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")
18 changes: 18 additions & 0 deletions python/tests/test_mimetypes_ext.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 5595562

Please sign in to comment.