Allow raw dicts in StateStore.set_encryption_info

Author: tulir

mautrix/client/state_store/abstract.py

@@ -5,7 +5,7 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from __future__ import annotations
 
-from typing import Awaitable
+from typing import Any, Awaitable
 from abc import ABC, abstractmethod
 
 from mautrix.types import (
@@ -135,7 +135,7 @@ async def get_encryption_info(self, room_id: RoomID) -> RoomEncryptionStateEvent
 
     @abstractmethod
     async def set_encryption_info(
-        self, room_id: RoomID, content: RoomEncryptionStateEventContent
+        self, room_id: RoomID, content: RoomEncryptionStateEventContent | dict[str, any]
     ) -> None:
         pass
 

mautrix/client/state_store/asyncpg/store.py

@@ -5,7 +5,7 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from __future__ import annotations
 
-from typing import NamedTuple
+from typing import Any, NamedTuple
 
 from mautrix.types import (
     Member,
@@ -14,6 +14,7 @@
     PowerLevelStateEventContent,
     RoomEncryptionStateEventContent,
     RoomID,
+    Serializable,
     UserID,
 )
 from mautrix.util.async_db import Database, Scheme
@@ -242,10 +243,12 @@ async def get_encryption_info(self, room_id: RoomID) -> RoomEncryptionStateEvent
         return RoomEncryptionStateEventContent.parse_json(row["encryption"])
 
     async def set_encryption_info(
-        self, room_id: RoomID, content: RoomEncryptionStateEventContent
+        self, room_id: RoomID, content: RoomEncryptionStateEventContent | dict[str, Any]
     ) -> None:
         q = (
             "INSERT INTO mx_room_state (room_id, is_encrypted, encryption) VALUES ($1, true, $2) "
             "ON CONFLICT (room_id) DO UPDATE SET is_encrypted=true, encryption=$2"
         )
-        await self.db.execute(q, room_id, content.json())
+        await self.db.execute(
+            q, room_id, content.json() if isinstance(content, Serializable) else content
+        )

mautrix/client/state_store/file.py

@@ -5,7 +5,7 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from __future__ import annotations
 
-from typing import IO
+from typing import IO, Any
 from pathlib import Path
 
 from mautrix.types import (
@@ -55,7 +55,7 @@ async def set_members(
         self._time_limited_flush()
 
     async def set_encryption_info(
-        self, room_id: RoomID, content: RoomEncryptionStateEventContent
+        self, room_id: RoomID, content: RoomEncryptionStateEventContent | dict[str, Any]
     ) -> None:
         await super().set_encryption_info(room_id, content)
         self._time_limited_flush()

mautrix/client/state_store/memory.py

@@ -187,6 +187,8 @@ async def get_encryption_info(self, room_id: RoomID) -> RoomEncryptionStateEvent
         return self.encryption.get(room_id)
 
     async def set_encryption_info(
-        self, room_id: RoomID, content: RoomEncryptionStateEventContent
+        self, room_id: RoomID, content: RoomEncryptionStateEventContent | dict[str, Any]
     ) -> None:
+        if not isinstance(content, RoomEncryptionStateEventContent):
+            content = RoomEncryptionStateEventContent.deserialize(content)
         self.encryption[room_id] = content

mautrix/client/state_store/sqlalchemy/mx_room_state.py

@@ -32,7 +32,7 @@ def python_type(self) -> Type[Serializable]:
 
     def process_bind_param(self, value: Serializable, dialect) -> str | None:
         if value is not None:
-            return json.dumps(value.serialize())
+            return json.dumps(value.serialize() if isinstance(value, Serializable) else value)
         return None
 
     def process_result_value(self, value: str, dialect) -> Serializable | None:

mautrix/client/state_store/sqlalchemy/sqlstatestore.py

@@ -5,6 +5,8 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from __future__ import annotations
 
+from typing import Any
+
 from mautrix.types import (
     Member,
     Membership,
@@ -174,7 +176,7 @@ async def get_encryption_info(self, room_id: RoomID) -> RoomEncryptionStateEvent
         return room.encryption
 
     async def set_encryption_info(
-        self, room_id: RoomID, content: RoomEncryptionStateEventContent
+        self, room_id: RoomID, content: RoomEncryptionStateEventContent | dict[str, Any]
     ) -> None:
         if not content:
             raise ValueError("content is empty")