Add missing type hints and fix most type errors except for Optionals.

mautrix_telegram/sqlstatestore.py

@@ -20,37 +20,39 @@ from sqlalchemy import orm
 
 from mautrix_appservice import StateStore
 
+from .types import MatrixUserId, MatrixRoomId
 from . import puppet as pu
 from .db import RoomState, UserProfile
 
 
 class SQLStateStore(StateStore):
-    def __init__(self, db) -> None:
+    def __init__(self, db: orm.Session) -> None:
         super().__init__()
         self.db = db  # type: orm.Session
         self.profile_cache = {}  # type: Dict[Tuple[str, str], UserProfile]
         self.room_state_cache = {}  # type: Dict[str, RoomState]
 
     @staticmethod
-    def is_registered(user: str) -> bool:
+    def is_registered(user: MatrixUserId) -> bool:
         puppet = pu.Puppet.get_by_mxid(user)
         return puppet.is_registered if puppet else False
 
     @staticmethod
-    def registered(user: str) -> None:
+    def registered(user: MatrixUserId) -> None:
         puppet = pu.Puppet.get_by_mxid(user)
         if puppet:
             puppet.is_registered = True
             puppet.save()
 
-    def update_state(self, event: dict) -> None:
+    def update_state(self, event: Dict) -> None:
         event_type = event["type"]
         if event_type == "m.room.power_levels":
             self.set_power_levels(event["room_id"], event["content"])
         elif event_type == "m.room.member":
             self.set_member(event["room_id"], event["state_key"], event["content"])
 
-    def _get_user_profile(self, room_id: str, user_id: str, create: bool = True) -> UserProfile:
+    def _get_user_profile(self, room_id: MatrixRoomId, user_id: MatrixUserId, create: bool = True
+                          ) -> UserProfile:
         key = (room_id, user_id)
         try:
             return self.profile_cache[key]
@@ -67,22 +69,22 @@ class SQLStateStore(StateStore):
             self.profile_cache[key] = profile
         return profile
 
-    def get_member(self, room: str, user: str) -> dict:
+    def get_member(self, room: MatrixRoomId, user: MatrixUserId) -> Dict:
         return self._get_user_profile(room, user).dict()
 
-    def set_member(self, room: str, user: str, member: dict) -> None:
+    def set_member(self, room: MatrixRoomId, user: MatrixUserId, member: Dict) -> None:
         profile = self._get_user_profile(room, user)
         profile.membership = member.get("membership", profile.membership or "leave")
         profile.displayname = member.get("displayname", profile.displayname)
         profile.avatar_url = member.get("avatar_url", profile.avatar_url)
         self.db.commit()
 
-    def set_membership(self, room: str, user: str, membership: str) -> None:
+    def set_membership(self, room: MatrixRoomId, user: MatrixUserId, membership: str) -> None:
         self.set_member(room, user, {
             "membership": membership,
         })
 
-    def _get_room_state(self, room_id: str, create: bool = True) -> RoomState:
+    def _get_room_state(self, room_id: MatrixRoomId, create: bool = True) -> RoomState:
         try:
             return self.room_state_cache[room_id]
         except KeyError:
@@ -96,13 +98,13 @@ class SQLStateStore(StateStore):
             self.room_state_cache[room_id] = room
         return room
 
-    def has_power_levels(self, room: str) -> bool:
+    def has_power_levels(self, room: MatrixRoomId) -> bool:
         return self._get_room_state(room).has_power_levels
 
-    def get_power_levels(self, room: str) -> dict:
+    def get_power_levels(self, room: MatrixRoomId) -> Dict:
         return self._get_room_state(room).power_levels
 
-    def set_power_level(self, room: str, user: str, level: int) -> None:
+    def set_power_level(self, room: MatrixRoomId, user: MatrixUserId, level: int) -> None:
         room_state = self._get_room_state(room)
         power_levels = room_state.power_levels
         if not power_levels:
@@ -114,7 +116,7 @@ class SQLStateStore(StateStore):
         room_state.power_levels = power_levels
         self.db.commit()
 
-    def set_power_levels(self, room: str, content: dict) -> None:
+    def set_power_levels(self, room: MatrixRoomId, content: Dict) -> None:
         state = self._get_room_state(room)
         state.power_levels = content
         self.db.commit()