Add helper to guess stream format (#308)

src/rokuecp/helpers.py

@@ -1,12 +1,67 @@
 """Helpers for Roku Client."""
 from __future__ import annotations
 
+import mimetypes
 from ipaddress import ip_address
 from socket import gaierror as SocketGIAError
 
+import yarl
+
 from .exceptions import RokuConnectionError
 from .resolver import ThreadedResolver
 
+MIME_TO_STREAM_FORMAT = {
+    "application/dash+xml": "dash",
+    "application/x-mpegURL": "hls",
+    "audio/mpeg": "mp3",
+    "audio/x-ms-wma": "wma",
+    "video/mp4": "mp4",
+    "video/quicktime": "mp4",
+    "video/x-matroska": "mkv",
+}
+
+
+def guess_stream_format(url: str, mime_type: str | None = None) -> str | None:
+    """Guess the Roku stream format for a given URL and MIME type.
+
+    Args:
+        url: The URL to determine stream format for.
+        mime_type: The MIME type to aid in stream format determination.
+
+    Returns:
+        The stream format or None if unable to determine stream format.
+    """
+    parsed = yarl.URL(url)
+    parsed_name = parsed.name.lower()
+
+    if mime_type is None:
+        mime_type, _ = mimetypes.guess_type(parsed.path)
+
+    if mime_type == "audio/mpeg" and parsed_name.endswith(".m4a"):
+        return "m4a"
+
+    if mime_type is None:
+        if parsed_name.endswith(".dash"):
+            return "dash"
+        if parsed_name.endswith(".mpd"):
+            return "dash"
+        if parsed_name.endswith(".m4v"):
+            return "mp4"
+        if parsed_name.endswith(".mks"):
+            return "mks"
+        if parsed_name.endswith(".mka"):
+            return "mka"
+        if ".ism/manifest" in parsed.path.lower():
+            return "ism"
+
+    if mime_type is None:
+        return None
+
+    if mime_type not in MIME_TO_STREAM_FORMAT:
+        return None
+
+    return MIME_TO_STREAM_FORMAT[mime_type]
+
 
 def is_ip_address(host: str) -> bool:
     """Determine if host is an IP Address.

tests/test_helpers.py

@@ -4,14 +4,36 @@
 
 import pytest
 from rokuecp.exceptions import RokuConnectionError
-from rokuecp.helpers import is_ip_address, resolve_hostname
+from rokuecp.helpers import guess_stream_format, is_ip_address, resolve_hostname
 
 from . import fake_addrinfo_results
 
 HOSTNAME = "roku.local"
 HOST = "192.168.1.2"
 
 
+def test_guess_stream_format() -> None:
+    """Test the guess_stream_format helper."""
+    assert guess_stream_format("/path/media.media") is None
+    assert guess_stream_format("/path/media.txt") is None
+
+    assert guess_stream_format("/path/media.mp4", "video/mp4") == "mp4"
+    assert guess_stream_format("/path/media.mp4") == "mp4"
+    assert guess_stream_format("/path/media.m4v") == "mp4"
+    assert guess_stream_format("/path/media.mov") == "mp4"
+    assert guess_stream_format("/path/media.mkv") == "mkv"
+    assert guess_stream_format("/path/media.mks") == "mks"
+    assert guess_stream_format("/path/media.m3u8") == "hls"
+    assert guess_stream_format("/path/media.dash") == "dash"
+    assert guess_stream_format("/path/media.mpd") == "dash"
+    assert guess_stream_format("/path/media.ism/manifest") == "ism"
+
+    assert guess_stream_format("/path/media.mp3") == "mp3"
+    assert guess_stream_format("/path/media.m4a") == "m4a"
+    assert guess_stream_format("/path/media.mka") == "mka"
+    assert guess_stream_format("/path/media.wma") == "wma"
+
+
 def test_is_ip_address() -> None:
     """Test the is_ip_address helper."""
     assert is_ip_address(HOST)