PYTHON-3788 add types to server_selectors.py (#1278)

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]
     )