Improved push typing (#11409)

Co-authored-by: Sean Quah <8349537+squahtx@users.noreply.github.com>

1 files changed, 42 insertions, 30 deletions

diff --git a/synapse/push/mailer.py b/synapse/push/mailer.py
index ce299ba3da..ba4f866487 100644
--- a/synapse/push/mailer.py
+++ b/synapse/push/mailer.py
@@ -14,7 +14,7 @@
 
 import logging
 import urllib.parse
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, TypeVar
+from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, TypeVar
 
 import bleach
 import jinja2
@@ -28,6 +28,14 @@ from synapse.push.presentable_names import (
     descriptor_from_member_events,
     name_from_member_event,
 )
+from synapse.push.push_types import (
+    EmailReason,
+    MessageVars,
+    NotifVars,
+    RoomVars,
+    TemplateVars,
+)
+from synapse.storage.databases.main.event_push_actions import EmailPushAction
 from synapse.storage.state import StateFilter
 from synapse.types import StateMap, UserID
 from synapse.util.async_helpers import concurrently_execute
@@ -135,7 +143,7 @@ class Mailer:
             % urllib.parse.urlencode(params)
         )
 
-        template_vars = {"link": link}
+        template_vars: TemplateVars = {"link": link}
 
         await self.send_email(
             email_address,
@@ -165,7 +173,7 @@ class Mailer:
             % urllib.parse.urlencode(params)
         )
 
-        template_vars = {"link": link}
+        template_vars: TemplateVars = {"link": link}
 
         await self.send_email(
             email_address,
@@ -196,7 +204,7 @@ class Mailer:
             % urllib.parse.urlencode(params)
         )
 
-        template_vars = {"link": link}
+        template_vars: TemplateVars = {"link": link}
 
         await self.send_email(
             email_address,
@@ -210,8 +218,8 @@ class Mailer:
         app_id: str,
         user_id: str,
         email_address: str,
-        push_actions: Iterable[Dict[str, Any]],
-        reason: Dict[str, Any],
+        push_actions: Iterable[EmailPushAction],
+        reason: EmailReason,
     ) -> None:
         """
         Send email regarding a user's room notifications
@@ -230,7 +238,7 @@ class Mailer:
             [pa["event_id"] for pa in push_actions]
         )
 
-        notifs_by_room: Dict[str, List[Dict[str, Any]]] = {}
+        notifs_by_room: Dict[str, List[EmailPushAction]] = {}
         for pa in push_actions:
             notifs_by_room.setdefault(pa["room_id"], []).append(pa)
 
@@ -258,7 +266,7 @@ class Mailer:
         # actually sort our so-called rooms_in_order list, most recent room first
         rooms_in_order.sort(key=lambda r: -(notifs_by_room[r][-1]["received_ts"] or 0))
 
-        rooms: List[Dict[str, Any]] = []
+        rooms: List[RoomVars] = []
 
         for r in rooms_in_order:
             roomvars = await self._get_room_vars(
@@ -289,7 +297,7 @@ class Mailer:
                 notifs_by_room, state_by_room, notif_events, reason
             )
 
-        template_vars = {
+        template_vars: TemplateVars = {
             "user_display_name": user_display_name,
             "unsubscribe_link": self._make_unsubscribe_link(
                 user_id, app_id, email_address
@@ -302,10 +310,10 @@ class Mailer:
         await self.send_email(email_address, summary_text, template_vars)
 
     async def send_email(
-        self, email_address: str, subject: str, extra_template_vars: Dict[str, Any]
+        self, email_address: str, subject: str, extra_template_vars: TemplateVars
     ) -> None:
         """Send an email with the given information and template text"""
-        template_vars = {
+        template_vars: TemplateVars = {
             "app_name": self.app_name,
             "server_name": self.hs.config.server.server_name,
         }
@@ -327,10 +335,10 @@ class Mailer:
         self,
         room_id: str,
         user_id: str,
-        notifs: Iterable[Dict[str, Any]],
+        notifs: Iterable[EmailPushAction],
         notif_events: Dict[str, EventBase],
         room_state_ids: StateMap[str],
-    ) -> Dict[str, Any]:
+    ) -> RoomVars:
         """
         Generate the variables for notifications on a per-room basis.
 
@@ -356,7 +364,7 @@ class Mailer:
 
         room_name = await calculate_room_name(self.store, room_state_ids, user_id)
 
-        room_vars: Dict[str, Any] = {
+        room_vars: RoomVars = {
             "title": room_name,
             "hash": string_ordinal_total(room_id),  # See sender avatar hash
             "notifs": [],
@@ -417,11 +425,11 @@ class Mailer:
 
     async def _get_notif_vars(
         self,
-        notif: Dict[str, Any],
+        notif: EmailPushAction,
         user_id: str,
         notif_event: EventBase,
         room_state_ids: StateMap[str],
-    ) -> Dict[str, Any]:
+    ) -> NotifVars:
         """
         Generate the variables for a single notification.
 
@@ -442,7 +450,7 @@ class Mailer:
             after_limit=CONTEXT_AFTER,
         )
 
-        ret = {
+        ret: NotifVars = {
             "link": self._make_notif_link(notif),
             "ts": notif["received_ts"],
             "messages": [],
@@ -461,8 +469,8 @@ class Mailer:
         return ret
 
     async def _get_message_vars(
-        self, notif: Dict[str, Any], event: EventBase, room_state_ids: StateMap[str]
-    ) -> Optional[Dict[str, Any]]:
+        self, notif: EmailPushAction, event: EventBase, room_state_ids: StateMap[str]
+    ) -> Optional[MessageVars]:
         """
         Generate the variables for a single event, if possible.
 
@@ -494,7 +502,9 @@ class Mailer:
 
         if sender_state_event:
             sender_name = name_from_member_event(sender_state_event)
-            sender_avatar_url = sender_state_event.content.get("avatar_url")
+            sender_avatar_url: Optional[str] = sender_state_event.content.get(
+                "avatar_url"
+            )
         else:
             # No state could be found, fallback to the MXID.
             sender_name = event.sender
@@ -504,7 +514,7 @@ class Mailer:
         # sender_hash % the number of default images to choose from
         sender_hash = string_ordinal_total(event.sender)
 
-        ret = {
+        ret: MessageVars = {
             "event_type": event.type,
             "is_historical": event.event_id != notif["event_id"],
             "id": event.event_id,
@@ -519,6 +529,8 @@ class Mailer:
             return ret
 
         msgtype = event.content.get("msgtype")
+        if not isinstance(msgtype, str):
+            msgtype = None
 
         ret["msgtype"] = msgtype
 
@@ -533,7 +545,7 @@ class Mailer:
         return ret
 
     def _add_text_message_vars(
-        self, messagevars: Dict[str, Any], event: EventBase
+        self, messagevars: MessageVars, event: EventBase
     ) -> None:
         """
         Potentially add a sanitised message body to the message variables.
@@ -543,8 +555,8 @@ class Mailer:
             event: The event under consideration.
         """
         msgformat = event.content.get("format")
-
-        messagevars["format"] = msgformat
+        if not isinstance(msgformat, str):
+            msgformat = None
 
         formatted_body = event.content.get("formatted_body")
         body = event.content.get("body")
@@ -555,7 +567,7 @@ class Mailer:
             messagevars["body_text_html"] = safe_text(body)
 
     def _add_image_message_vars(
-        self, messagevars: Dict[str, Any], event: EventBase
+        self, messagevars: MessageVars, event: EventBase
     ) -> None:
         """
         Potentially add an image URL to the message variables.
@@ -570,7 +582,7 @@ class Mailer:
     async def _make_summary_text_single_room(
         self,
         room_id: str,
-        notifs: List[Dict[str, Any]],
+        notifs: List[EmailPushAction],
         room_state_ids: StateMap[str],
         notif_events: Dict[str, EventBase],
         user_id: str,
@@ -685,10 +697,10 @@ class Mailer:
 
     async def _make_summary_text(
         self,
-        notifs_by_room: Dict[str, List[Dict[str, Any]]],
+        notifs_by_room: Dict[str, List[EmailPushAction]],
         room_state_ids: Dict[str, StateMap[str]],
         notif_events: Dict[str, EventBase],
-        reason: Dict[str, Any],
+        reason: EmailReason,
     ) -> str:
         """
         Make a summary text for the email when multiple rooms have notifications.
@@ -718,7 +730,7 @@ class Mailer:
     async def _make_summary_text_from_member_events(
         self,
         room_id: str,
-        notifs: List[Dict[str, Any]],
+        notifs: List[EmailPushAction],
         room_state_ids: StateMap[str],
         notif_events: Dict[str, EventBase],
     ) -> str:
@@ -805,7 +817,7 @@ class Mailer:
             base_url = "https://matrix.to/#"
         return "%s/%s" % (base_url, room_id)
 
-    def _make_notif_link(self, notif: Dict[str, str]) -> str:
+    def _make_notif_link(self, notif: EmailPushAction) -> str:
         """
         Generate a link to open an event in the web client.