Annotations for user_erasure_store (matrix-org#11313)

David Robertson · web-flow · commit c99da2d0794d · 2021-11-11T19:22:19.000Z
I'm not sure why this was excluded---it seemed to be passing for me. But
it's easy enough to fixup.
diff --git a/changelog.d/11313.misc b/changelog.d/11313.misc
@@ -0,0 +1 @@
+Add type hints to storage classes.
diff --git a/mypy.ini b/mypy.ini
@@ -53,7 +53,6 @@ exclude = (?x)
    |synapse/storage/databases/main/stats.py
    |synapse/storage/databases/main/transactions.py
    |synapse/storage/databases/main/user_directory.py
-   |synapse/storage/databases/main/user_erasure_store.py
    |synapse/storage/schema/
 
    |tests/api/test_auth.py
@@ -184,6 +183,9 @@ disallow_untyped_defs = True
 [mypy-synapse.storage.databases.main.room_batch]
 disallow_untyped_defs = True
 
+[mypy-synapse.storage.databases.main.user_erasure_store]
+disallow_untyped_defs = True
+
 [mypy-synapse.storage.util.*]
 disallow_untyped_defs = True
 
diff --git a/synapse/storage/databases/main/user_erasure_store.py b/synapse/storage/databases/main/user_erasure_store.py
@@ -14,11 +14,12 @@
 
 from typing import Dict, Iterable
 
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.database import LoggingTransaction
+from synapse.storage.databases.main import CacheInvalidationWorkerStore
 from synapse.util.caches.descriptors import cached, cachedList
 
 
-class UserErasureWorkerStore(SQLBaseStore):
+class UserErasureWorkerStore(CacheInvalidationWorkerStore):
     @cached()
     async def is_user_erased(self, user_id: str) -> bool:
         """
@@ -69,7 +70,7 @@ async def mark_user_erased(self, user_id: str) -> None:
             user_id: full user_id to be erased
         """
 
-        def f(txn):
+        def f(txn: LoggingTransaction) -> None:
             # first check if they are already in the list
             txn.execute("SELECT 1 FROM erased_users WHERE user_id = ?", (user_id,))
             if txn.fetchone():
@@ -89,7 +90,7 @@ async def mark_user_not_erased(self, user_id: str) -> None:
             user_id: full user_id to be un-erased
         """
 
-        def f(txn):
+        def f(txn: LoggingTransaction) -> None:
             # first check if they are already in the list
             txn.execute("SELECT 1 FROM erased_users WHERE user_id = ?", (user_id,))
             if not txn.fetchone():
