Expose an enum with the different content types

src/h_vialib/__init__.py

@@ -1,4 +1,4 @@
 """Library functions for Via related products."""
 
-from h_vialib.client import ViaClient, ViaDoc
+from h_vialib.client import ContentType, ViaClient, ViaDoc
 from h_vialib.configuration import Configuration

src/h_vialib/client.py

@@ -1,13 +1,20 @@
 """Helper classes for clients using Via proxying."""
 
 import re
+from enum import Enum
+from typing import Optional
 from urllib.parse import parse_qsl, urlencode, urlparse
 
 from webob.multidict import MultiDict
 
 from h_vialib.secure import Encryption, ViaSecureURL
 
 
+class ContentType(str, Enum):
+    PDF = "pdf"
+    HTML = "html"
+
+
 class ViaDoc:
     """A doc we want to proxy with content type."""
 
@@ -20,19 +27,9 @@ def __init__(self, url, content_type=None):
         self.url = url
 
         if content_type is None and self._GOOGLE_DRIVE_REGEX.match(url):
-            content_type = "pdf"
-
-        self._content_type = content_type
-
-    @property
-    def is_html(self):
-        """Check if document is known to be a HTML."""
-        return self._content_type == "html"
+            content_type = ContentType.PDF
 
-    @property
-    def is_pdf(self):
-        """Check if document is known to be a pdf."""
-        return self._content_type == "pdf"
+        self.content_type = content_type
 
 
 class ViaClient:  # pylint: disable=too-few-public-methods
@@ -59,7 +56,12 @@ def __init__(self, secret, service_url=None, html_service_url=None):
 
     # pylint:disable=too-many-arguments
     def url_for(
-        self, url, content_type=None, options=None, blocked_for=None, headers=None
+        self,
+        url,
+        content_type: Optional[ContentType] = None,
+        options=None,
+        blocked_for=None,
+        headers=None,
     ):
         """Generate a Via URL to display a given URL.
 
@@ -88,7 +90,7 @@ def url_for(
         if blocked_for:
             query["via.blocked_for"] = blocked_for
 
-        if doc.is_html:
+        if doc.content_type == ContentType.HTML:
             # Optimisation to skip routing for documents we know are HTML
             via_url = self._url_for_html(doc.url, query)
         else:
@@ -101,7 +103,7 @@ def _url_for(self, doc, query):
             raise ValueError("Cannot rewrite URLs without a service URL")
 
         # Optimisation to skip routing for documents we know are PDFs
-        path = "/pdf" if doc.is_pdf else "/route"
+        path = "/pdf" if doc.content_type == ContentType.PDF else "/route"
 
         query["url"] = doc.url
 

tests/unit/h_vialib/client_test.py

@@ -1,37 +1,24 @@
 import pytest
 from h_matchers import Any
 
-from h_vialib import ViaClient, ViaDoc
+from h_vialib import ContentType, ViaClient, ViaDoc
 
 
 class TestViaDoc:
     @pytest.mark.parametrize(
-        "url,content_type,is_pdf",
+        "url,content_type,expected_content_type",
         (
-            ("http://example.com", None, False),
-            ("http://example.com", "html", False),
-            ("http://example.com", "pdf", True),
+            ("http://example.com", None, None),
+            ("http://example.com", ContentType.HTML, ContentType.HTML),
+            ("http://example.com", ContentType.PDF, ContentType.PDF),
             # We know about Google Drive links and assume them to be PDF
-            ("https://drive.google.com/uc?id=0&export=download", None, True),
+            ("https://drive.google.com/uc?id=0&export=download", None, ContentType.PDF),
         ),
     )
-    def test_is_pdf(self, url, content_type, is_pdf):
+    def test_content_type(self, url, content_type, expected_content_type):
         doc = ViaDoc(url, content_type)
 
-        assert doc.is_pdf == is_pdf
-
-    @pytest.mark.parametrize(
-        "url,content_type,is_html",
-        (
-            ("http://example.com", None, False),
-            ("http://example.com", "pdf", False),
-            ("http://example.com", "html", True),
-        ),
-    )
-    def test_is_htmlf(self, url, content_type, is_html):
-        doc = ViaDoc(url, content_type)
-
-        assert doc.is_html == is_html
+        assert doc.content_type == expected_content_type
 
 
 class TestViaClient: