From f563676c097b830346acc7a4ce3e910c6b10c4c3 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Wed, 6 Oct 2021 18:55:25 +0100
Subject: [PATCH] `disallow-untyped-defs` for `synapse.state` (#11004)

* `disallow-untyped-defs` for `synapse.state`

Much smaller than I was expecting!
---
 changelog.d/11004.misc    | 1 +
 mypy.ini                  | 3 +++
 synapse/state/__init__.py | 2 +-
 synapse/state/v1.py       | 4 ++--
 synapse/state/v2.py       | 2 +-
 5 files changed, 8 insertions(+), 4 deletions(-)
 create mode 100644 changelog.d/11004.misc

diff --git a/changelog.d/11004.misc b/changelog.d/11004.misc
new file mode 100644
index 0000000000..821033710a
--- /dev/null
+++ b/changelog.d/11004.misc
@@ -0,0 +1 @@
+Add further type hints to `synapse.state`.
\ No newline at end of file
diff --git a/mypy.ini b/mypy.ini
index 86459bdcb6..a052d49c71 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -99,6 +99,9 @@ disallow_untyped_defs = True
 [mypy-synapse.rest.*]
 disallow_untyped_defs = True
 
+[mypy-synapse.state.*]
+disallow_untyped_defs = True
+
 [mypy-synapse.util.batching_queue]
 disallow_untyped_defs = True
 
diff --git a/synapse/state/__init__.py b/synapse/state/__init__.py
index c981df3f18..5cf2e12575 100644
--- a/synapse/state/__init__.py
+++ b/synapse/state/__init__.py
@@ -118,7 +118,7 @@ class _StateCacheEntry:
         else:
             self.state_id = _gen_state_id()
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.state)
 
 
diff --git a/synapse/state/v1.py b/synapse/state/v1.py
index 017e6fd92d..ffe6207a3c 100644
--- a/synapse/state/v1.py
+++ b/synapse/state/v1.py
@@ -225,7 +225,7 @@ def _resolve_with_state(
     conflicted_state_ids: StateMap[Set[str]],
     auth_event_ids: StateMap[str],
     state_map: Dict[str, EventBase],
-):
+) -> MutableStateMap[str]:
     conflicted_state = {}
     for key, event_ids in conflicted_state_ids.items():
         events = [state_map[ev_id] for ev_id in event_ids if ev_id in state_map]
@@ -362,7 +362,7 @@ def _resolve_normal_events(
 
 
 def _ordered_events(events: Iterable[EventBase]) -> List[EventBase]:
-    def key_func(e):
+    def key_func(e: EventBase) -> Tuple[int, str]:
         # we have to use utf-8 rather than ascii here because it turns out we allow
         # people to send us events with non-ascii event IDs :/
         return -int(e.depth), hashlib.sha1(e.event_id.encode("utf-8")).hexdigest()
diff --git a/synapse/state/v2.py b/synapse/state/v2.py
index 586b0e12fe..bd18eefd58 100644
--- a/synapse/state/v2.py
+++ b/synapse/state/v2.py
@@ -481,7 +481,7 @@ async def _reverse_topological_power_sort(
         if idx % _AWAIT_AFTER_ITERATIONS == 0:
             await clock.sleep(0)
 
-    def _get_power_order(event_id):
+    def _get_power_order(event_id: str) -> Tuple[int, int, str]:
         ev = event_map[event_id]
         pl = event_to_pl[event_id]