From cdb89dcefe9f4d7035f898cd77cd514fa69c2673 Mon Sep 17 00:00:00 2001
From: Patrick Cloke <clokep@users.noreply.github.com>
Date: Thu, 28 Sep 2023 07:01:46 -0400
Subject: [PATCH] Improve state types. (#16395)

---
 changelog.d/16395.misc |  1 +
 synapse/state/v2.py    |  5 ++---
 tests/state/test_v2.py | 13 ++++++++-----
 3 files changed, 11 insertions(+), 8 deletions(-)
 create mode 100644 changelog.d/16395.misc

diff --git a/changelog.d/16395.misc b/changelog.d/16395.misc
new file mode 100644
index 0000000000..93ceaeafc9
--- /dev/null
+++ b/changelog.d/16395.misc
@@ -0,0 +1 @@
+Improve type hints.
diff --git a/synapse/state/v2.py b/synapse/state/v2.py
index 1752f95db8..b2e63aed1e 100644
--- a/synapse/state/v2.py
+++ b/synapse/state/v2.py
@@ -23,7 +23,6 @@ from typing import (
     Generator,
     Iterable,
     List,
-    Mapping,
     Optional,
     Sequence,
     Set,
@@ -269,7 +268,7 @@ async def _get_power_level_for_sender(
 
 async def _get_auth_chain_difference(
     room_id: str,
-    state_sets: Sequence[Mapping[Any, str]],
+    state_sets: Sequence[StateMap[str]],
     unpersisted_events: Dict[str, EventBase],
     state_res_store: StateResolutionStore,
 ) -> Set[str]:
@@ -405,7 +404,7 @@ def _seperate(
 
     # mypy doesn't understand that discarding None above means that conflicted
     # state is StateMap[Set[str]], not StateMap[Set[Optional[Str]]].
-    return unconflicted_state, conflicted_state  # type: ignore
+    return unconflicted_state, conflicted_state  # type: ignore[return-value]
 
 
 def _is_power_event(event: EventBase) -> bool:
diff --git a/tests/state/test_v2.py b/tests/state/test_v2.py
index 2e3f2318d9..6a2f7584f6 100644
--- a/tests/state/test_v2.py
+++ b/tests/state/test_v2.py
@@ -719,7 +719,10 @@ class AuthChainDifferenceTestCase(unittest.TestCase):
         persisted_events = {a.event_id: a, b.event_id: b}
         unpersited_events = {c.event_id: c}
 
-        state_sets = [{"a": a.event_id, "b": b.event_id}, {"c": c.event_id}]
+        state_sets = [
+            {("a", ""): a.event_id, ("b", ""): b.event_id},
+            {("c", ""): c.event_id},
+        ]
 
         store = TestStateResolutionStore(persisted_events)
 
@@ -774,8 +777,8 @@ class AuthChainDifferenceTestCase(unittest.TestCase):
         unpersited_events = {c.event_id: c, d.event_id: d}
 
         state_sets = [
-            {"a": a.event_id, "b": b.event_id},
-            {"c": c.event_id, "d": d.event_id},
+            {("a", ""): a.event_id, ("b", ""): b.event_id},
+            {("c", ""): c.event_id, ("d", ""): d.event_id},
         ]
 
         store = TestStateResolutionStore(persisted_events)
@@ -841,8 +844,8 @@ class AuthChainDifferenceTestCase(unittest.TestCase):
         unpersited_events = {c.event_id: c, d.event_id: d, e.event_id: e}
 
         state_sets = [
-            {"a": a.event_id, "b": b.event_id, "e": e.event_id},
-            {"c": c.event_id, "d": d.event_id},
+            {("a", ""): a.event_id, ("b", ""): b.event_id, ("e", ""): e.event_id},
+            {("c", ""): c.event_id, ("d", ""): d.event_id},
         ]
 
         store = TestStateResolutionStore(persisted_events)