service: Split GrpcChannelPool into a separate submodule (#376)

* service: Split GrpcChannelPool into a separate submodule

* service: Use typing_extensions.Self

* service: Fix lint warnings

Author: bkeryan

ni_measurementlink_service/_channelpool.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import sys
+from threading import Lock
+from types import TracebackType
+from typing import (
+    Dict,
+    Literal,
+    Optional,
+    Type,
+    TYPE_CHECKING,
+)
+
+import grpc
+
+from ni_measurementlink_service._loggers import ClientLogger
+
+if TYPE_CHECKING:
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
+
+class GrpcChannelPool(object):
+    """Class that manages gRPC channel lifetimes."""
+
+    def __init__(self) -> None:
+        """Initialize the GrpcChannelPool object."""
+        self._lock: Lock = Lock()
+        self._channel_cache: Dict[str, grpc.Channel] = {}
+
+    def __enter__(self: Self) -> Self:
+        """Enter the runtime context of the GrpcChannelPool."""
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Literal[False]:
+        """Exit the runtime context of the GrpcChannelPool."""
+        self.close()
+        return False
+
+    def get_channel(self, target: str) -> grpc.Channel:
+        """Return a gRPC channel.
+
+        Args:
+            target (str): The server address
+
+        """
+        new_channel = None
+        with self._lock:
+            if target not in self._channel_cache:
+                self._lock.release()
+                new_channel = grpc.insecure_channel(target)
+                if ClientLogger.is_enabled():
+                    new_channel = grpc.intercept_channel(new_channel, ClientLogger())
+                self._lock.acquire()
+                if target not in self._channel_cache:
+                    self._channel_cache[target] = new_channel
+                    new_channel = None
+            channel = self._channel_cache[target]
+
+        # Close new_channel if it was not stored in _channel_cache.
+        if new_channel is not None:
+            new_channel.close()
+
+        return channel
+
+    def close(self) -> None:
+        """Close channels opened by get_channel()."""
+        with self._lock:
+            for channel in self._channel_cache.values():
+                channel.close()
+            self._channel_cache.clear()

ni_measurementlink_service/measurement/service.py

@@ -7,7 +7,6 @@
 from enum import Enum, EnumMeta
 from os import path
 from pathlib import Path
-from threading import Lock
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
@@ -18,21 +17,22 @@
     Literal,
     Optional,
     Type,
-    TypeVar,
     Union,
 )
 
 import grpc
 from google.protobuf.descriptor import EnumDescriptor
 
 from ni_measurementlink_service import _datatypeinfo
+from ni_measurementlink_service._channelpool import (  # re-export
+    GrpcChannelPool as GrpcChannelPool,
+)
 from ni_measurementlink_service._internal import grpc_servicer
 from ni_measurementlink_service._internal.discovery_client import DiscoveryClient
 from ni_measurementlink_service._internal.parameter import (
     metadata as parameter_metadata,
 )
 from ni_measurementlink_service._internal.service_manager import GrpcService
-from ni_measurementlink_service._loggers import ClientLogger
 from ni_measurementlink_service.measurement.info import (
     DataType,
     MeasurementInfo,
@@ -49,6 +49,11 @@
     else:
         from typing_extensions import TypeGuard
 
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
     SupportedEnumType = Union[Type[Enum], _EnumTypeWrapper]
 
 
@@ -83,67 +88,6 @@ def abort(self, code: grpc.StatusCode, details: str) -> None:
         grpc_servicer.measurement_service_context.get().abort(code, details)
 
 
-# Eventually, these can be replaced with typing.Self (Python >= 3.11).
-_TGrpcChannelPool = TypeVar("_TGrpcChannelPool", bound="GrpcChannelPool")
-_TMeasurementService = TypeVar("_TMeasurementService", bound="MeasurementService")
-
-
-class GrpcChannelPool(object):
-    """Class that manages gRPC channel lifetimes."""
-
-    def __init__(self) -> None:
-        """Initialize the GrpcChannelPool object."""
-        self._lock: Lock = Lock()
-        self._channel_cache: Dict[str, grpc.Channel] = {}
-
-    def __enter__(self: _TGrpcChannelPool) -> _TGrpcChannelPool:
-        """Enter the runtime context of the GrpcChannelPool."""
-        return self
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        traceback: Optional[TracebackType],
-    ) -> Literal[False]:
-        """Exit the runtime context of the GrpcChannelPool."""
-        self.close()
-        return False
-
-    def get_channel(self, target: str) -> grpc.Channel:
-        """Return a gRPC channel.
-
-        Args:
-            target (str): The server address
-
-        """
-        new_channel = None
-        with self._lock:
-            if target not in self._channel_cache:
-                self._lock.release()
-                new_channel = grpc.insecure_channel(target)
-                if ClientLogger.is_enabled():
-                    new_channel = grpc.intercept_channel(new_channel, ClientLogger())
-                self._lock.acquire()
-                if target not in self._channel_cache:
-                    self._channel_cache[target] = new_channel
-                    new_channel = None
-            channel = self._channel_cache[target]
-
-        # Close new_channel if it was not stored in _channel_cache.
-        if new_channel is not None:
-            new_channel.close()
-
-        return channel
-
-    def close(self) -> None:
-        """Close channels opened by get_channel()."""
-        with self._lock:
-            for channel in self._channel_cache.values():
-                channel.close()
-            self._channel_cache.clear()
-
-
 class MeasurementService:
     """Class that supports registering and hosting a python function as a gRPC service.
 
@@ -451,7 +395,7 @@ def close_service(self) -> None:
         self.grpc_service.stop()
         self.channel_pool.close()
 
-    def __enter__(self: _TMeasurementService) -> _TMeasurementService:
+    def __enter__(self: Self) -> Self:
         """Enter the runtime context related to the measurement service."""
         return self
 

ni_measurementlink_service/session_management.py

@@ -2,10 +2,12 @@
 from __future__ import annotations
 
 import abc
+import sys
 import warnings
 from functools import cached_property
 from types import TracebackType
 from typing import (
+    TYPE_CHECKING,
     Any,
     Iterable,
     List,
@@ -14,7 +16,6 @@
     Optional,
     Sequence,
     Type,
-    TypeVar,
 )
 
 import grpc
@@ -29,6 +30,12 @@
     session_management_service_pb2_grpc,
 )
 
+if TYPE_CHECKING:
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
 GRPC_SERVICE_INTERFACE_NAME = "ni.measurementlink.sessionmanagement.v1.SessionManagementService"
 GRPC_SERVICE_CLASS = "ni.measurementlink.sessionmanagement.v1.SessionManagementService"
 
@@ -175,10 +182,6 @@ def _convert_session_info_to_grpc(
     )
 
 
-# Eventually, this can be replaced with typing.Self (Python >= 3.11).
-_TReservation = TypeVar("_TReservation", bound="BaseReservation")
-
-
 class BaseReservation(abc.ABC):
     """Manages session reservation."""
 
@@ -191,7 +194,7 @@ def __init__(
         self._session_manager = session_manager
         self._session_info = session_info
 
-    def __enter__(self: _TReservation) -> _TReservation:
+    def __enter__(self: Self) -> Self:
         """Context management protocol. Returns self."""
         return self