Rebase onto upstream (a4d95fd)
#12
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import mimetypes
|
||||
from enum import Enum
|
||||
|
||||
from loguru import logger
|
||||
|
||||
@@ -13,6 +14,14 @@ except ImportError:
|
||||
HAS_MAGIC = False
|
||||
|
||||
|
||||
class MediaKind(Enum):
|
||||
"""Media type classification."""
|
||||
IMAGE = "image"
|
||||
VIDEO = "video"
|
||||
AUDIO = "audio"
|
||||
DOCUMENT = "document"
|
||||
|
||||
|
||||
def detect_mime(path: str, content: bytes | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of media file.
|
||||
@@ -47,3 +56,23 @@ def detect_mime(path: str, content: bytes | None = None) -> str:
|
||||
|
||||
# Fallback
|
||||
return "application/octet-stream"
|
||||
|
||||
|
||||
def classify_media(mime: str) -> MediaKind:
|
||||
"""
|
||||
Classify MIME type into media kind.
|
||||
|
||||
Args:
|
||||
mime: MIME type string (e.g., "image/jpeg")
|
||||
|
||||
Returns:
|
||||
MediaKind enum value
|
||||
"""
|
||||
if mime.startswith("image/"):
|
||||
return MediaKind.IMAGE
|
||||
if mime.startswith("video/"):
|
||||
return MediaKind.VIDEO
|
||||
if mime.startswith("audio/"):
|
||||
return MediaKind.AUDIO
|
||||
# Everything else is a document
|
||||
return MediaKind.DOCUMENT
|
||||
|
||||
@@ -65,3 +65,37 @@ def test_detect_mime_malformed_content():
|
||||
# Should fallback to extension detection, not crash
|
||||
mime = detect_mime("test.mp4", malformed)
|
||||
assert mime == "video/mp4"
|
||||
|
||||
|
||||
def test_classify_media_image():
|
||||
"""Test classification of image MIME types."""
|
||||
from nanobot.channels.telegram_media import MediaKind, classify_media
|
||||
|
||||
assert classify_media("image/jpeg") == MediaKind.IMAGE
|
||||
assert classify_media("image/png") == MediaKind.IMAGE
|
||||
assert classify_media("image/webp") == MediaKind.IMAGE
|
||||
|
||||
|
||||
def test_classify_media_video():
|
||||
"""Test classification of video MIME types."""
|
||||
from nanobot.channels.telegram_media import MediaKind, classify_media
|
||||
|
||||
assert classify_media("video/mp4") == MediaKind.VIDEO
|
||||
assert classify_media("video/quicktime") == MediaKind.VIDEO
|
||||
|
||||
|
||||
def test_classify_media_audio():
|
||||
"""Test classification of audio MIME types."""
|
||||
from nanobot.channels.telegram_media import MediaKind, classify_media
|
||||
|
||||
assert classify_media("audio/mpeg") == MediaKind.AUDIO
|
||||
assert classify_media("audio/ogg") == MediaKind.AUDIO
|
||||
|
||||
|
||||
def test_classify_media_document():
|
||||
"""Test classification of document MIME types."""
|
||||
from nanobot.channels.telegram_media import MediaKind, classify_media
|
||||
|
||||
assert classify_media("application/pdf") == MediaKind.DOCUMENT
|
||||
assert classify_media("text/plain") == MediaKind.DOCUMENT
|
||||
assert classify_media("application/octet-stream") == MediaKind.DOCUMENT
|
||||
|
||||
Reference in New Issue
Block a user