diff --git a/pymongo/server_selectors.py b/pymongo/server_selectors.py index aa9d26b5f..9a6701557 100644 --- a/pymongo/server_selectors.py +++ b/pymongo/server_selectors.py @@ -13,15 +13,25 @@ # permissions and limitations under the License. """Criteria to select some ServerDescriptions from a TopologyDescription.""" +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Optional, TypeVar from pymongo.server_type import SERVER_TYPE +if TYPE_CHECKING: + from pymongo.server_description import ServerDescription + from pymongo.topology_description import TopologyDescription + + +T = TypeVar("T") + class Selection: """Input or output of a server selector function.""" @classmethod - def from_topology_description(cls, topology_description): + def from_topology_description(cls, topology_description: TopologyDescription) -> Selection: known_servers = topology_description.known_servers primary = None for sd in known_servers: @@ -36,54 +46,60 @@ class Selection: primary, ) - def __init__(self, topology_description, server_descriptions, common_wire_version, primary): + def __init__( + self, + topology_description: TopologyDescription, + server_descriptions: List[ServerDescription], + common_wire_version: Optional[int], + primary: Optional[ServerDescription], + ): self.topology_description = topology_description self.server_descriptions = server_descriptions self.primary = primary self.common_wire_version = common_wire_version - def with_server_descriptions(self, server_descriptions): + def with_server_descriptions(self, server_descriptions: List[ServerDescription]) -> Selection: return Selection( self.topology_description, server_descriptions, self.common_wire_version, self.primary ) - def secondary_with_max_last_write_date(self): + def secondary_with_max_last_write_date(self) -> Optional[ServerDescription]: secondaries = secondary_server_selector(self) if secondaries.server_descriptions: return max(secondaries.server_descriptions, key=lambda sd: sd.last_write_date) return None @property - def primary_selection(self): + def primary_selection(self) -> Selection: primaries = [self.primary] if self.primary else [] return self.with_server_descriptions(primaries) @property - def heartbeat_frequency(self): + def heartbeat_frequency(self) -> int: return self.topology_description.heartbeat_frequency @property - def topology_type(self): + def topology_type(self) -> int: return self.topology_description.topology_type - def __bool__(self): + def __bool__(self) -> bool: return bool(self.server_descriptions) - def __getitem__(self, item): + def __getitem__(self, item: int) -> ServerDescription: return self.server_descriptions[item] -def any_server_selector(selection): +def any_server_selector(selection: T) -> T: return selection -def readable_server_selector(selection): +def readable_server_selector(selection: Selection) -> Selection: return selection.with_server_descriptions( [s for s in selection.server_descriptions if s.is_readable] ) -def writable_server_selector(selection): +def writable_server_selector(selection: Selection) -> Selection: return selection.with_server_descriptions( [s for s in selection.server_descriptions if s.is_writable] )