From 8dc666f785c653e958728a74758bdb8afb494de4 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Thu, 11 Nov 2021 16:49:28 +0000
Subject: [PATCH] Correct type hint for room_batch.py (#11310)

Co-authored-by: Patrick Cloke <clokep@users.noreply.github.com>
---
 changelog.d/11310.misc                       |  1 +
 mypy.ini                                     |  4 +++-
 synapse/storage/databases/main/room_batch.py | 18 ++++++++----------
 3 files changed, 12 insertions(+), 11 deletions(-)
 create mode 100644 changelog.d/11310.misc

diff --git a/changelog.d/11310.misc b/changelog.d/11310.misc
new file mode 100644
index 000000000..e5c12445d
--- /dev/null
+++ b/changelog.d/11310.misc
@@ -0,0 +1 @@
+Add type hints to storage classes.
\ No newline at end of file
diff --git a/mypy.ini b/mypy.ini
index a06c7fc66..48dfdfa0e 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -46,7 +46,6 @@ exclude = (?x)
    |synapse/storage/databases/main/push_rule.py
    |synapse/storage/databases/main/receipts.py
    |synapse/storage/databases/main/room.py
-   |synapse/storage/databases/main/room_batch.py
    |synapse/storage/databases/main/roommember.py
    |synapse/storage/databases/main/search.py
    |synapse/storage/databases/main/signatures.py
@@ -183,6 +182,9 @@ disallow_untyped_defs = True
 [mypy-synapse.storage.databases.main.client_ips]
 disallow_untyped_defs = True
 
+[mypy-synapse.storage.databases.main.room_batch]
+disallow_untyped_defs = True
+
 [mypy-synapse.storage.util.*]
 disallow_untyped_defs = True
 
diff --git a/synapse/storage/databases/main/room_batch.py b/synapse/storage/databases/main/room_batch.py
index 97b261843..39e80f6f5 100644
--- a/synapse/storage/databases/main/room_batch.py
+++ b/synapse/storage/databases/main/room_batch.py
@@ -39,13 +39,11 @@ class RoomBatchStore(SQLBaseStore):
 
     async def store_state_group_id_for_event_id(
         self, event_id: str, state_group_id: int
-    ) -> Optional[str]:
-        {
-            await self.db_pool.simple_upsert(
-                table="event_to_state_groups",
-                keyvalues={"event_id": event_id},
-                values={"state_group": state_group_id, "event_id": event_id},
-                # Unique constraint on event_id so we don't have to lock
-                lock=False,
-            )
-        }
+    ) -> None:
+        await self.db_pool.simple_upsert(
+            table="event_to_state_groups",
+            keyvalues={"event_id": event_id},
+            values={"state_group": state_group_id, "event_id": event_id},
+            # Unique constraint on event_id so we don't have to lock
+            lock=False,
+        )