Add annotations to metadata interceptor, other style improvements.

Change-Id: I58040150b1dbbbbfe4d2e19dc6384d8f26d9f6c1

Author: BenRKarl

google/ads/googleads/interceptors/__init__.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 
 from .helpers import mask_message
-from .interceptor import Interceptor, MetadataType
+from .interceptor import Interceptor, MetadataType, ContinuationType
 from .metadata_interceptor import MetadataInterceptor
 from .exception_interceptor import ExceptionInterceptor
 from .logging_interceptor import LoggingInterceptor
 
 __all__ = [
+    "ContinuationType",
     "ExceptionInterceptor",
     "Interceptor",
     "LoggingInterceptor",

google/ads/googleads/interceptors/exception_interceptor.py

@@ -19,12 +19,12 @@
 so it translates the error to a GoogleAdsFailure instance and raises it.
 """
 
-from typing import Callable, Optional, Union
+from typing import Optional, Union
 
 from google.protobuf.message import Message
 import grpc
 
-from google.ads.googleads.interceptors import Interceptor
+from google.ads.googleads.interceptors import Interceptor, ContinuationType
 from google.ads.googleads.response_wrappers import _UnaryStreamWrapper, _UnaryUnaryWrapper
 
 
@@ -73,7 +73,7 @@ def _handle_grpc_failure(self, response: Union[grpc.Call, grpc.Future]):
 
     def intercept_unary_unary(
         self,
-        continuation: Callable[[grpc.ClientCallDetails, Message], grpc.Call],
+        continuation: ContinuationType,
         client_call_details: grpc.ClientCallDetails,
         request: Message,
     ):
@@ -98,7 +98,7 @@ def intercept_unary_unary(
                 indicative of a GoogleAdsException, or if the exception has a
                 status code of INTERNAL or RESOURCE_EXHAUSTED.
         """
-        response: grpc._interceptor._UnaryOutcome = continuation(client_call_details, request)
+        response: grpc.Call = continuation(client_call_details, request)
         exception: Optional[grpc.RpcError] = response.exception()
 
         if exception:
@@ -110,7 +110,7 @@ def intercept_unary_unary(
 
     def intercept_unary_stream(
         self,
-        continuation: Callable[[grpc.ClientCallDetails, Message], grpc.Call],
+        continuation: ContinuationType,
         client_call_details: grpc.ClientCallDetails,
         request: Message,
     ):
@@ -135,7 +135,7 @@ def intercept_unary_stream(
                 indicative of a GoogleAdsException, or if the exception has a
                 status code of INTERNAL or RESOURCE_EXHAUSTED.
         """
-        response: grpc._interceptor._UnaryOutcome = continuation(client_call_details, request)
+        response: grpc.Call = continuation(client_call_details, request)
         return _UnaryStreamWrapper(
             response,
             self._handle_grpc_failure,

google/ads/googleads/interceptors/interceptor.py

@@ -25,36 +25,38 @@
 from typing import (
     Any,
     AnyStr,
+    Callable,
     Dict,
     Optional,
     Sequence,
     Tuple
 )
 
 from google.protobuf.message import DecodeError, Message
-from grpc import Call, ClientCallDetails, StatusCode, CallCredentials, RpcError
+import grpc
+# from grpc import Call, ClientCallDetails, StatusCode, CallCredentials, RpcError
 
 from google.ads.googleads.errors import GoogleAdsException
 
 
 _REQUEST_ID_KEY: str = "request-id"
 # Codes that are retried upon by google.api_core.
-_RETRY_STATUS_CODES: Tuple[StatusCode]= (StatusCode.INTERNAL, StatusCode.RESOURCE_EXHAUSTED)
+_RETRY_STATUS_CODES: Tuple[grpc.StatusCode]= (grpc.StatusCode.INTERNAL, grpc.StatusCode.RESOURCE_EXHAUSTED)
 
 MetadataType = Sequence[Tuple[str, AnyStr]]
-
+ContinuationType = Callable[[grpc.ClientCallDetails, Message], grpc.Call]
 
 class Interceptor:
     _SENSITIVE_INFO_MASK: str = "REDACTED"
 
     @dataclass
-    class _ClientCallDetails(ClientCallDetails):
+    class _ClientCallDetails(grpc.ClientCallDetails):
         """Wrapper class for initializing a new ClientCallDetails instance."""
 
         method: str
         timeout: Optional[float]
         metadata: Optional[MetadataType]
-        credentials: Optional[CallCredentials]
+        credentials: Optional[grpc.CallCredentials]
 
     @classmethod
     def get_request_id_from_metadata(
@@ -134,7 +136,7 @@ def default_serializer(value: Any) -> Optional[str]:
 
     @classmethod
     def get_trailing_metadata_from_interceptor_exception(
-        cls, exception: RpcError
+        cls, exception: grpc.RpcError
     ) -> MetadataType:
         """Retrieves trailing metadata from an exception object.
 
@@ -164,8 +166,8 @@ def get_client_call_details_instance(
         method: str,
         timeout: float,
         metadata: MetadataType,
-        credentials: Optional[CallCredentials] = None,
-    ) -> ClientCallDetails:
+        credentials: Optional[grpc.CallCredentials] = None,
+    ) -> grpc.ClientCallDetails:
         """Initializes an instance of the ClientCallDetails with the given data.
 
         Args:
@@ -186,7 +188,7 @@ def __init__(self, api_version: str):
         )
         self._api_version: str = api_version
 
-    def _get_error_from_response(self, response: Call) -> Call:
+    def _get_error_from_response(self, response: grpc.Call) -> grpc.Call:
         """Attempts to wrap failed responses as GoogleAdsException instances.
 
         Handles failed gRPC responses of by attempting to convert them
@@ -210,8 +212,8 @@ def _get_error_from_response(self, response: Call) -> Call:
             Exception: If not a GoogleAdsException or RpcException the error
                 will be raised as-is.
         """
-        status_code: StatusCode = response.code()
-        response_exception: Call = response.exception()
+        status_code: grpc.StatusCode = response.code()
+        response_exception: grpc.Call = response.exception()
 
         if status_code not in _RETRY_STATUS_CODES:
             trailing_metadata: MetadataType = response.trailing_metadata()

google/ads/googleads/interceptors/logging_interceptor.py

@@ -21,12 +21,14 @@
 
 import json
 import logging
-from typing import Callable, Optional, Union
+from typing import Optional, Union
 
 from google.protobuf.message import Message as ProtobufMessageType
 import grpc
 
-from google.ads.googleads.interceptors import Interceptor, MetadataType, mask_message
+from google.ads.googleads.interceptors import (
+    Interceptor, MetadataType, ContinuationType, mask_message
+)
 
 
 class LoggingInterceptor(
@@ -339,7 +341,7 @@ def log_request(
 
     def intercept_unary_unary(
         self,
-        continuation: Callable[[grpc.ClientCallDetails, ProtobufMessageType], grpc.Call],
+        continuation: ContinuationType,
         client_call_details: grpc.ClientCallDetails,
         request: ProtobufMessageType,
     ) -> Union[grpc.Call, grpc.Future]:
@@ -357,7 +359,7 @@ def intercept_unary_unary(
         Returns:
             A grpc.Call/grpc.Future instance representing a service response.
         """
-        response: grpc._interceptor._UnaryOutcome = continuation(client_call_details, request)
+        response: grpc.Call = continuation(client_call_details, request)
 
         if self.logger.isEnabledFor(logging.WARNING):
             self.log_request(client_call_details, request, response)
@@ -366,7 +368,7 @@ def intercept_unary_unary(
 
     def intercept_unary_stream(
         self,
-        continuation: Callable[[grpc.ClientCallDetails, ProtobufMessageType], grpc.Call],
+        continuation: ContinuationType,
         client_call_details: grpc.ClientCallDetails,
         request: ProtobufMessageType,
     ) -> Union[grpc.Call, grpc.Future]:
@@ -388,7 +390,7 @@ def on_rpc_complete(response_future: grpc.Future) -> None:
             if self.logger.isEnabledFor(logging.WARNING):
                 self.log_request(client_call_details, request, response_future)
 
-        response: grpc._interceptor._UnaryOutcome = continuation(client_call_details, request)
+        response: grpc.Call = continuation(client_call_details, request)
 
         response.add_done_callback(on_rpc_complete)
         self._cache = response.get_cache()

google/ads/googleads/interceptors/metadata_interceptor.py

@@ -24,18 +24,21 @@
 # request user-agent directly by the google-api-core package:
 # https://github.com/googleapis/python-api-core/issues/416
 from importlib import metadata
+from typing import List, Optional, Tuple, Union
 
+from google.protobuf.internal import api_implementation
+from google.protobuf.message import Message as ProtobufMessageType
+import grpc
+
+from google.ads.googleads.interceptors import Interceptor, MetadataType, ContinuationType
+
+
+# Determine which version of the package is installed.
 try:
     _PROTOBUF_VERSION = metadata.version("protobuf")
 except metadata.PackageNotFoundError:
     _PROTOBUF_VERSION = None
 
-
-from google.protobuf.internal import api_implementation
-from grpc import UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor
-
-from .interceptor import Interceptor
-
 # Determine which protobuf implementation is being used.
 if api_implementation.Type() == "cpp":
     _PB_IMPL_HEADER = "+c"
@@ -44,18 +47,19 @@
 else:
     _PB_IMPL_HEADER = ""
 
-
 class MetadataInterceptor(
-    Interceptor, UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor
+    Interceptor,
+    grpc.UnaryUnaryClientInterceptor,
+    grpc.UnaryStreamClientInterceptor
 ):
     """An interceptor that appends custom metadata to requests."""
 
     def __init__(
         self,
-        developer_token,
-        login_customer_id,
-        linked_customer_id=None,
-        use_cloud_org_for_api_access=None,
+        developer_token: str,
+        login_customer_id: Optional[str]= None,
+        linked_customer_id: Optional[str] = None,
+        use_cloud_org_for_api_access: Optional[bool] = None,
     ):
         """Initialization method for this class.
 
@@ -69,21 +73,26 @@ def __init__(
                 levels. Use this flag only if you are enrolled into a limited
                 pilot that supports this configuration
         """
-        self.developer_token_meta = ("developer-token", developer_token)
-        self.login_customer_id_meta = (
+        self.developer_token_meta: Tuple[str, str] = (
+            "developer-token",
+            developer_token,
+        )
+        self.login_customer_id_meta: Optional[Tuple[str, str]] = (
             ("login-customer-id", login_customer_id)
             if login_customer_id
             else None
         )
-        self.linked_customer_id_meta = (
+        self.linked_customer_id_meta: Optional[Tuple[str, str]] = (
             ("linked-customer-id", linked_customer_id)
             if linked_customer_id
             else None
         )
-        self.use_cloud_org_for_api_access = use_cloud_org_for_api_access
+        self.use_cloud_org_for_api_access: Optional[
+            bool
+        ] = use_cloud_org_for_api_access
 
     def _update_client_call_details_metadata(
-        self, client_call_details, metadata
+        self, client_call_details: grpc.ClientCallDetails, metadata: MetadataType
     ):
         """Updates the client call details with additional metadata.
 
@@ -95,7 +104,7 @@ def _update_client_call_details_metadata(
             An new instance of grpc.ClientCallDetails with additional metadata
             from the GoogleAdsClient.
         """
-        client_call_details = self.get_client_call_details_instance(
+        client_call_details: grpc.ClientCallDetails = self.get_client_call_details_instance(
             client_call_details.method,
             client_call_details.timeout,
             metadata,
@@ -104,7 +113,12 @@ def _update_client_call_details_metadata(
 
         return client_call_details
 
-    def _intercept(self, continuation, client_call_details, request):
+    def _intercept(
+        self,
+        continuation: ContinuationType,
+        client_call_details: grpc.ClientCallDetails,
+        request: ProtobufMessageType,
+    ) -> Union[grpc.Call, grpc.Future]:
         """Generic interceptor used for Unary-Unary and Unary-Stream requests.
 
         Args:
@@ -118,9 +132,9 @@ def _intercept(self, continuation, client_call_details, request):
             A grpc.Call/grpc.Future instance representing a service response.
         """
         if client_call_details.metadata is None:
-            metadata = []
+            metadata: MetadataType = []
         else:
-            metadata = list(client_call_details.metadata)
+            metadata: MetadataType = list(client_call_details.metadata)
 
         # If self.use_cloud_org_for_api_access is not True, add the developer
         # token to the request's metadata
@@ -135,31 +149,36 @@ def _intercept(self, continuation, client_call_details, request):
 
         # TODO: This logic should be updated or removed once the following is
         # fixed: https://github.com/googleapis/python-api-core/issues/416
-        for i, metadatum in enumerate(metadata):
+        for i, metadatum_tuple in enumerate(metadata):
             # Check if the user agent header key is in the current metadatum
-            if "x-goog-api-client" in metadatum and _PROTOBUF_VERSION:
+            if "x-goog-api-client" in metadatum_tuple and _PROTOBUF_VERSION:
                 # Convert the tuple to a list so it can be modified.
-                metadatum = list(metadatum)
+                metadatum: List[str] = list(metadatum_tuple)
                 # Check that "pb" isn't already included in the user agent.
                 if "pb" not in metadatum[1]:
                     # Append the protobuf version key value pair to the end of
                     # the string.
                     metadatum[1] += f" pb/{_PROTOBUF_VERSION}{_PB_IMPL_HEADER}"
                     # Convert the metadatum back to a tuple.
-                    metadatum = tuple(metadatum)
+                    metadatum_tuple: Tuple[str, str] = tuple(metadatum)
                     # Splice the metadatum back in its original position in
                     # order to preserve the order of the metadata list.
-                    metadata[i] = metadatum
+                    metadata[i] = metadatum_tuple
                     # Exit the loop since we already found the user agent.
                     break
 
-        client_call_details = self._update_client_call_details_metadata(
+        client_call_details: grpc.ClientCallDetails = self._update_client_call_details_metadata(
             client_call_details, metadata
         )
 
         return continuation(client_call_details, request)
 
-    def intercept_unary_unary(self, continuation, client_call_details, request):
+    def intercept_unary_unary(
+        self,
+        continuation: ContinuationType,
+        client_call_details: grpc.ClientCallDetails,
+        request: ProtobufMessageType,
+    ) -> Union[grpc.Call, grpc.Future]:
         """Intercepts and appends custom metadata for Unary-Unary requests.
 
         Overrides abstract method defined in grpc.UnaryUnaryClientInterceptor.
@@ -177,8 +196,11 @@ def intercept_unary_unary(self, continuation, client_call_details, request):
         return self._intercept(continuation, client_call_details, request)
 
     def intercept_unary_stream(
-        self, continuation, client_call_details, request
-    ):
+        self,
+        continuation: ContinuationType,
+        client_call_details: grpc.ClientCallDetails,
+        request: ProtobufMessageType,
+    ) -> Union[grpc.Call, grpc.Future]:
         """Intercepts and appends custom metadata to Unary-Stream requests.
 
         Overrides abstract method defined in grpc.UnaryStreamClientInterceptor.