Cleanup of components.

1 files changed, 33 insertions, 11 deletions

diff --git a/mu4web/components.py b/mu4web/components.py
index a441390..e9701de 100644
--- a/mu4web/components.py
+++ b/mu4web/components.py
@@ -1,8 +1,12 @@
 """Various HTML "components"."""
 
-from email.headerregistry import Address
+from email.headerregistry import (
+    Address,
+    BaseHeader,
+    UniqueAddressHeader,
+)
 from .html_render import HTML
-from typing import Any, cast, Optional
+from typing import cast, Optional
 from urllib.parse import urlencode
 from email.message import EmailMessage
 
@@ -18,19 +22,25 @@ def mailto(addr: str) -> HTML:
     return ('a', {'href': f'mailto:{addr}'}, addr)
 
 
-def header_format(key: str, value: Any) -> HTML:
-    """Format email headers to HTML."""
+def header_format(key: str, value: BaseHeader) -> HTML:
+    """
+    Format email headers to HTML.
+
+    https://docs.python.org/3/library/email.headerregistry.html#email.headerregistry.HeaderRegistry
+    """
     if key in ['to', 'cc', 'bcc']:
+        assert isinstance(value, UniqueAddressHeader), \
+            f"Expected UniqueAddressHeader, got {type(value)}"
         return ('ul', *[('li', *format_email(addr))
                         for addr in value.addresses])
     elif key == 'from':
+        assert isinstance(value, UniqueAddressHeader), \
+            f"Expected UniqueAddressHeader, got {type(value)}"
         return format_email(value.addresses[0])
     elif key == 'in-reply-to':
-        # type(value) == email.headerregistry._UnstructuredHeader
         id = str(value).strip("<>")
         return ['<', ('a', {'href': '?' + urlencode({'id': id})}, id), '>']
     else:
-        # assert type(value) == str, f"Weird value in header {value!r}"
         return str(value)
 
 
@@ -81,8 +91,7 @@ def login_page(returnto: Optional[str] = None) -> HTML:
             ('label', {'for': 'remember'}, 'Kom ihåg mig')),
             ('input', {'type': 'hidden',
                        'name': 'returnto',
-                       'value': returnto})
-            if returnto else [],
+                       'value': returnto}) if returnto else [],
             ('input', {'type': 'submit', 'value': 'Logga in'}),
             )
 
@@ -104,9 +113,22 @@ def login_prompt() -> HTML:
 
 
 def flashed_messages(messages: list[str] | list[tuple[str, str]]) -> HTML:
-    """Return Flasks flashed messages, formatted as a list."""
-    return ('ul', {'class': 'flashes'},
-            *[('li', msg) for msg in messages])
+    """
+    Return Flasks flashed messages, formatted as a list.
+
+    :param messages:
+        Either a list of messages, or a list of (category, message)
+        messages. See `flask.get_flashed_messages` for more information.
+    """
+    if not messages:
+        return []
+    if type(messages[0]) == tuple:
+        messages = cast(list[tuple[str, str]], messages)
+        return ('ul', {'class': 'flashes'},
+                *[('li', msg) for (_, msg) in messages])
+    else:
+        return ('ul', {'class': 'flashes'},
+                *[('li', msg) for msg in messages])
 
 
 def include_stylesheet(path: str) -> HTML: