Add more type hints to the main state store. (#12267)

This commit is contained in:
Patrick Cloke 2022-03-31 13:38:09 -04:00 committed by GitHub
parent 5e88143dff
commit 11df4ec6c2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 16 additions and 10 deletions

1
changelog.d/12267.misc Normal file
View File

@ -0,0 +1 @@
Add missing type hints for storage.

View File

@ -43,7 +43,6 @@ exclude = (?x)
|synapse/storage/databases/main/event_federation.py |synapse/storage/databases/main/event_federation.py
|synapse/storage/databases/main/push_rule.py |synapse/storage/databases/main/push_rule.py
|synapse/storage/databases/main/roommember.py |synapse/storage/databases/main/roommember.py
|synapse/storage/databases/main/state.py
|synapse/storage/schema/ |synapse/storage/schema/
|tests/api/test_auth.py |tests/api/test_auth.py

View File

@ -12,9 +12,10 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import collections.abc
import logging import logging
from typing import TYPE_CHECKING, Collection, Iterable, Optional, Set, Tuple from typing import TYPE_CHECKING, Dict, Iterable, Optional, Set, Tuple
from frozendict import frozendict
from synapse.api.constants import EventTypes, Membership from synapse.api.constants import EventTypes, Membership
from synapse.api.errors import NotFoundError, UnsupportedRoomVersionError from synapse.api.errors import NotFoundError, UnsupportedRoomVersionError
@ -29,7 +30,7 @@ from synapse.storage.database import (
from synapse.storage.databases.main.events_worker import EventsWorkerStore from synapse.storage.databases.main.events_worker import EventsWorkerStore
from synapse.storage.databases.main.roommember import RoomMemberWorkerStore from synapse.storage.databases.main.roommember import RoomMemberWorkerStore
from synapse.storage.state import StateFilter from synapse.storage.state import StateFilter
from synapse.types import JsonDict, StateMap from synapse.types import JsonDict, JsonMapping, StateMap
from synapse.util.caches import intern_string from synapse.util.caches import intern_string
from synapse.util.caches.descriptors import cached, cachedList from synapse.util.caches.descriptors import cached, cachedList
@ -132,7 +133,7 @@ class StateGroupWorkerStore(EventsWorkerStore, SQLBaseStore):
return room_version return room_version
async def get_room_predecessor(self, room_id: str) -> Optional[dict]: async def get_room_predecessor(self, room_id: str) -> Optional[JsonMapping]:
"""Get the predecessor of an upgraded room if it exists. """Get the predecessor of an upgraded room if it exists.
Otherwise return None. Otherwise return None.
@ -158,9 +159,10 @@ class StateGroupWorkerStore(EventsWorkerStore, SQLBaseStore):
predecessor = create_event.content.get("predecessor", None) predecessor = create_event.content.get("predecessor", None)
# Ensure the key is a dictionary # Ensure the key is a dictionary
if not isinstance(predecessor, collections.abc.Mapping): if not isinstance(predecessor, (dict, frozendict)):
return None return None
# The keys must be strings since the data is JSON.
return predecessor return predecessor
async def get_create_event_for_room(self, room_id: str) -> EventBase: async def get_create_event_for_room(self, room_id: str) -> EventBase:
@ -306,7 +308,9 @@ class StateGroupWorkerStore(EventsWorkerStore, SQLBaseStore):
list_name="event_ids", list_name="event_ids",
num_args=1, num_args=1,
) )
async def _get_state_group_for_events(self, event_ids: Collection[str]) -> JsonDict: async def _get_state_group_for_events(
self, event_ids: Iterable[str]
) -> Dict[str, int]:
"""Returns mapping event_id -> state_group""" """Returns mapping event_id -> state_group"""
rows = await self.db_pool.simple_select_many_batch( rows = await self.db_pool.simple_select_many_batch(
table="event_to_state_groups", table="event_to_state_groups",
@ -521,7 +525,7 @@ class MainStateBackgroundUpdateStore(RoomMemberWorkerStore):
) )
for user_id in potentially_left_users - joined_users: for user_id in potentially_left_users - joined_users:
await self.mark_remote_user_device_list_as_unsubscribed(user_id) await self.mark_remote_user_device_list_as_unsubscribed(user_id) # type: ignore[attr-defined]
return batch_size return batch_size

View File

@ -17,7 +17,7 @@ import logging
import typing import typing
from enum import Enum, auto from enum import Enum, auto
from sys import intern from sys import intern
from typing import Any, Callable, Dict, List, Optional, Sized from typing import Any, Callable, Dict, List, Optional, Sized, TypeVar
import attr import attr
from prometheus_client.core import Gauge from prometheus_client.core import Gauge
@ -195,8 +195,10 @@ KNOWN_KEYS = {
) )
} }
T = TypeVar("T", Optional[str], str)
def intern_string(string: Optional[str]) -> Optional[str]:
def intern_string(string: T) -> T:
"""Takes a (potentially) unicode string and interns it if it's ascii""" """Takes a (potentially) unicode string and interns it if it's ascii"""
if string is None: if string is None:
return None return None