From ab3b3b7c01a496590a8ed9011e1f3393d813cedc Mon Sep 17 00:00:00 2001
From: "Bendik R. Brenne" <bendik@konstant.no>
Date: Mon, 16 Sep 2024 20:49:39 +0200
Subject: [PATCH] feat(api): add file info fetching and improve error handling

- Implement fetch_file_info method in NrkPodcastAPI
- Add NrkPsApiNotFoundError and NrkPsApiRateLimitError exceptions
- Enhance error handling in _request method
- Introduce fetch_file_info utility function with caching
---
 nrk_psapi/api.py        | 18 ++++++++++++++++--
 nrk_psapi/const.py      |  1 +
 nrk_psapi/exceptions.py |  4 ++++
 nrk_psapi/utils.py      | 31 ++++++++++++++++++++++++++++---
 4 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/nrk_psapi/api.py b/nrk_psapi/api.py
index 6c6f6e3..c96dc31 100644
--- a/nrk_psapi/api.py
+++ b/nrk_psapi/api.py
@@ -4,6 +4,7 @@
 
 import asyncio
 from dataclasses import dataclass
+from http import HTTPStatus
 import socket
 
 from aiohttp.client import ClientError, ClientResponseError, ClientSession
@@ -18,6 +19,8 @@
     NrkPsApiConnectionError,
     NrkPsApiConnectionTimeoutError,
     NrkPsApiError,
+    NrkPsApiNotFoundError,
+    NrkPsApiRateLimitError,
 )
 from .models.catalog import (
     Episode,
@@ -48,7 +51,7 @@
     SearchResultType,
     SingleLetter,
 )
-from .utils import get_nested_items, sanitize_string
+from .utils import fetch_file_info, get_nested_items, sanitize_string
 
 
 @dataclass
@@ -154,9 +157,16 @@ async def _request(
             raise NrkPsApiConnectionTimeoutError(
                 "Timeout occurred while connecting to NRK API"
             ) from exception
+        except ClientResponseError as exception:
+            if exception.status == HTTPStatus.TOO_MANY_REQUESTS:
+                raise NrkPsApiRateLimitError(
+                    "Too many requests to NRK API. Try again later."
+                ) from exception
+            if exception.status == HTTPStatus.NOT_FOUND:
+                raise NrkPsApiNotFoundError("Resource not found") from exception
+            raise NrkPsApiError from exception
         except (
             ClientError,
-            ClientResponseError,
             socket.gaierror,
         ) as exception:
             msg = "Error occurred while communicating with NRK API"
@@ -478,6 +488,10 @@ async def curated_podcasts(self) -> Curated:
                     )
         return Curated(sections=sections)
 
+    async def fetch_file_info(self, url: URL | str):
+        """Proxies call to `utils.fetch_file_info`, passing on self.session."""
+        return await fetch_file_info(url, self.session)
+
     async def close(self) -> None:
         """Close open client session."""
         if self.session and self._close_session:
diff --git a/nrk_psapi/const.py b/nrk_psapi/const.py
index aee1c3f..f57a793 100644
--- a/nrk_psapi/const.py
+++ b/nrk_psapi/const.py
@@ -1,4 +1,5 @@
 """NRK Podcast API constants."""
+
 import logging
 
 LOGGER = logging.getLogger(__name__)
diff --git a/nrk_psapi/exceptions.py b/nrk_psapi/exceptions.py
index dbe7d7b..e5164be 100644
--- a/nrk_psapi/exceptions.py
+++ b/nrk_psapi/exceptions.py
@@ -5,6 +5,10 @@ class NrkPsApiError(Exception):
     """Generic NrkPs exception."""
 
 
+class NrkPsApiNotFoundError(NrkPsApiError):
+    """NrkPs not found exception."""
+
+
 class NrkPsApiConnectionError(NrkPsApiError):
     """NrkPs connection exception."""
 
diff --git a/nrk_psapi/utils.py b/nrk_psapi/utils.py
index 9afd311..b65aab6 100644
--- a/nrk_psapi/utils.py
+++ b/nrk_psapi/utils.py
@@ -1,6 +1,13 @@
 from __future__ import annotations
 
 import re
+from typing import TYPE_CHECKING
+
+from aiohttp import ClientSession
+from asyncstdlib import cache
+
+if TYPE_CHECKING:
+    from yarl import URL
 
 
 def get_nested_items(data: dict[str, any], items_key: str) -> list[dict[str, any]]:
@@ -19,6 +26,24 @@ def get_nested_items(data: dict[str, any], items_key: str) -> list[dict[str, any
 def sanitize_string(s: str):
     """Sanitize a string to be used as a URL parameter."""
 
-    s = s.lower().replace(' ', '_')
-    s = s.replace('æ', 'ae').replace('ø', 'oe').replace('å', 'aa')
-    return re.sub(r'^[0-9_]+', '', re.sub(r'[^a-z0-9_]', '', s))[:50].rstrip('_')
+    s = s.lower().replace(" ", "_")
+    s = s.replace("æ", "ae").replace("ø", "oe").replace("å", "aa")
+    return re.sub(r"^[0-9_]+", "", re.sub(r"[^a-z0-9_]", "", s))[:50].rstrip("_")
+
+
+@cache
+async def fetch_file_info(
+    url: URL | str, session: ClientSession | None = None
+) -> tuple[int, str]:
+    """Retrieve content-length and content-type for the given URL."""
+    close_session = False
+    if session is None:
+        session = ClientSession()
+        close_session = True
+
+    response = await session.head(url)
+    content_length = response.headers.get("Content-Length")
+    mime_type = response.headers.get("Content-Type")
+    if close_session:
+        await session.close()
+    return int(content_length), mime_type