From 5595562fbacb9984d6af4ab21dee0a5affbafc56 Mon Sep 17 00:00:00 2001
From: Aron Carroll <aron@replicate.com>
Date: Fri, 18 Oct 2024 15:54:23 +0100
Subject: [PATCH] Add support for image/webp to mimetypes package

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.
---
 python/cog/__init__.py             |  5 +++++
 python/cog/mimetypes_ext.py        | 21 +++++++++++++++++++++
 python/tests/test_mimetypes_ext.py | 18 ++++++++++++++++++
 3 files changed, 44 insertions(+)
 create mode 100644 python/cog/mimetypes_ext.py
 create mode 100644 python/tests/test_mimetypes_ext.py

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)