Fix issues found in linting.

2 files changed, 23 insertions, 13 deletions

diff --git a/synapse/handlers/sync.py b/synapse/handlers/sync.py
index 396deb39ce..32f4ed228a 100644
--- a/synapse/handlers/sync.py
+++ b/synapse/handlers/sync.py
@@ -1384,15 +1384,18 @@ class SyncHandler:
             # Non-lazy syncs should never include partially stated rooms.
             # Exclude all partially stated rooms from this sync.
             results = await self.store.is_partial_state_rooms(mutable_joined_room_ids)
-            for room_id, result in results.items():
-                mutable_rooms_to_exclude.add(room_id)
+            mutable_rooms_to_exclude.update(
+                room_id
+                for room_id, is_partial_state in results.items()
+                if is_partial_state
+            )
 
         # Incremental eager syncs should additionally include rooms that
         # - we are joined to
         # - are full-stated
         # - became fully-stated at some point during the sync period
         #   (These rooms will have been omitted during a previous eager sync.)
-        forced_newly_joined_room_ids = set()
+        forced_newly_joined_room_ids: Set[str] = set()
         if since_token and not sync_config.filter_collection.lazy_load_members():
             un_partial_stated_rooms = (
                 await self.store.get_un_partial_stated_rooms_between(
@@ -1402,8 +1405,11 @@ class SyncHandler:
                 )
             )
             results = await self.store.is_partial_state_rooms(un_partial_stated_rooms)
-            for room_id, result in results.items():
-                forced_newly_joined_room_ids.add(room_id)
+            forced_newly_joined_room_ids.update(
+                room_id
+                for room_id, is_partial_state in results.items()
+                if is_partial_state
+            )
 
         # Now we have our list of joined room IDs, exclude as configured and freeze
         joined_room_ids = frozenset(
diff --git a/synapse/storage/databases/main/room.py b/synapse/storage/databases/main/room.py
index aa7e2d6e55..f9c7be1b3e 100644
--- a/synapse/storage/databases/main/room.py
+++ b/synapse/storage/databases/main/room.py
@@ -1275,7 +1275,9 @@ class RoomWorkerStore(CacheInvalidationWorkerStore):
         return entry is not None
 
     @cachedList(cached_method_name="is_partial_State_room", list_name="room_ids")
-    async def is_partial_state_rooms(self, room_ids: StrCollection) -> Mapping[str, bool]:
+    async def is_partial_state_rooms(
+        self, room_ids: StrCollection
+    ) -> Mapping[str, bool]:
         """Checks if this room has partial state.
 
         Returns true if this is a "partial-state" room, which means that the state
@@ -1283,13 +1285,15 @@ class RoomWorkerStore(CacheInvalidationWorkerStore):
         complete.
         """
 
-        entries = set(await self.db_pool.simple_select_many_batch(
-            table="partial_state_rooms",
-            column="room_id",
-            iterable=room_ids,
-            retcols=("room_id",),
-            desc="is_partial_state_room",
-        ))
+        entries = set(
+            await self.db_pool.simple_select_many_batch(
+                table="partial_state_rooms",
+                column="room_id",
+                iterable=room_ids,
+                retcols=("room_id",),
+                desc="is_partial_state_room",
+            )
+        )
 
         return {room_id: room_id in entries for room_id in room_ids}