Fix type annotations in sequencers

plamut · Nov 30, 2021 · Nov 15, 2021 · Nov 15, 2021 · Nov 16, 2021 · Nov 16, 2021
commit d2790cc2c08eb207f277c6ad9bc09e227ceb9e11
@@ -22,15 +22,14 @@
 
 if typing.TYPE_CHECKING:  # pragma: NO COVER
     from concurrent import futures
-    from google.api_core import retry
+    from google.pubsub_v1.services.publisher.client import OptionalRetry
 
 
 class Sequencer(metaclass=abc.ABCMeta):
     """The base class for sequencers for Pub/Sub publishing. A sequencer
        sequences messages to be published.
     """
 
-    @staticmethod
     @abc.abstractmethod
     def is_finished(self) -> bool:  # pragma: NO COVER
         """ Whether the sequencer is finished and should be cleaned up.
@@ -40,7 +39,6 @@ def is_finished(self) -> bool:  # pragma: NO COVER
         """
         raise NotImplementedError
 
-    @staticmethod
     @abc.abstractmethod
     def unpause(self) -> None:  # pragma: NO COVER
         """ Unpauses this sequencer.
@@ -51,12 +49,11 @@ def unpause(self) -> None:  # pragma: NO COVER
         """
         raise NotImplementedError
 
-    @staticmethod
     @abc.abstractmethod
     def publish(
         self,
         message: gapic_types.PubsubMessage,
-        retry: "retry.Retry" = None,
+        retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
         timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
     ) -> "futures.Future":  # pragma: NO COVER
         """ Publish message for this ordering key.

@@ -14,21 +14,24 @@
 
 import enum
 import collections
-import concurrent.futures as futures
 import threading
 import typing
-from typing import Iterable, Sequence
+from typing import cast, Deque, Iterable, Optional, Sequence, Union
 
 from google.api_core import gapic_v1
+from google.cloud.pubsub_v1.publisher import futures
 from google.cloud.pubsub_v1.publisher import exceptions
 from google.cloud.pubsub_v1.publisher._sequencer import base as sequencer_base
 from google.cloud.pubsub_v1.publisher._batch import base as batch_base
 from google.pubsub_v1 import types as gapic_types
 
 if typing.TYPE_CHECKING:  # pragma: NO COVER
-    from google.api_core import retry
-    from google.cloud.pubsub_v1 import PublisherClient
     from google.cloud.pubsub_v1.publisher import _batch
+    from google.cloud.pubsub_v1.publisher.client import Client as PublisherClient
+    from google.pubsub_v1.services.publisher.client import OptionalRetry
+
+
+_TimeoutType = Union[gapic_types.TimeoutType, gapic_v1.method._MethodDefault]
 
 
 class _OrderedSequencerStatus(str, enum.Enum):
@@ -101,7 +104,7 @@ def __init__(self, client: "PublisherClient", topic: str, ordering_key: str):
         # Batches ordered from first (head/left) to last (right/tail).
         # Invariant: always has at least one batch after the first publish,
         # unless paused or stopped.
-        self._ordered_batches = collections.deque()
+        self._ordered_batches: Deque["_batch.thread.Batch"] = collections.deque()
         # See _OrderedSequencerStatus for valid state transitions.
         self._state = _OrderedSequencerStatus.ACCEPTING_MESSAGES
 
@@ -237,8 +240,8 @@ def unpause(self) -> None:
 
     def _create_batch(
         self,
-        commit_retry: "retry.Retry" = gapic_v1.method.DEFAULT,
-        commit_timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
+        commit_retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
+        commit_timeout: _TimeoutType = gapic_v1.method.DEFAULT,
     ) -> "_batch.thread.Batch":
         """ Create a new batch using the client's batch class and other stored
             settings.
@@ -262,8 +265,8 @@ def _create_batch(
     def publish(
         self,
         message: gapic_types.PubsubMessage,
-        retry: "retry.Retry" = gapic_v1.method.DEFAULT,
-        timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
+        retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
+        timeout: _TimeoutType = gapic_v1.method.DEFAULT,
     ) -> futures.Future:
         """ Publish message for this ordering key.
 
@@ -289,12 +292,12 @@ def publish(
         """
         with self._state_lock:
             if self._state == _OrderedSequencerStatus.PAUSED:
-                future = futures.Future()
+                future: Optional[futures.Future] = futures.Future()
                 exception = exceptions.PublishToPausedOrderingKeyException(
                     self._ordering_key
                 )
-                future.set_exception(exception)
-                return future
+                cast(futures.Future, future).set_exception(exception)
+                return cast(futures.Future, future)
 
             # If waiting to be cleaned-up, convert to accepting messages to
             # prevent this sequencer from being cleaned-up only to have another

@@ -13,17 +13,21 @@
 # limitations under the License.
 
 import typing
+from typing import Optional, Union
 
 from google.api_core import gapic_v1
 
 from google.cloud.pubsub_v1.publisher._sequencer import base
 from google.pubsub_v1 import types as gapic_types
 
 if typing.TYPE_CHECKING:  # pragma: NO COVER
-    from concurrent import futures
-    from google.api_core import retry
-    from google.cloud.pubsub_v1 import PublisherClient
     from google.cloud.pubsub_v1.publisher import _batch
+    from google.cloud.pubsub_v1.publisher import futures
+    from google.cloud.pubsub_v1.publisher.client import Client as PublisherClient
+    from google.pubsub_v1.services.publisher.client import OptionalRetry
+
+
+_TimeoutType = Union[gapic_types.TimeoutType, gapic_v1.method._MethodDefault]
 
 
 class UnorderedSequencer(base.Sequencer):
@@ -35,7 +39,7 @@ class UnorderedSequencer(base.Sequencer):
     def __init__(self, client: "PublisherClient", topic: str):
         self._client = client
         self._topic = topic
-        self._current_batch = None
+        self._current_batch: Optional["_batch.thread.Batch"] = None
         self._stopped = False
 
     def is_finished(self) -> bool:
@@ -88,8 +92,8 @@ def unpause(self) -> typing.NoReturn:
 
     def _create_batch(
         self,
-        commit_retry: "retry.Retry" = gapic_v1.method.DEFAULT,
-        commit_timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
+        commit_retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
+        commit_timeout: _TimeoutType = gapic_v1.method.DEFAULT,
     ) -> "_batch.thread.Batch":
         """ Create a new batch using the client's batch class and other stored
             settings.
@@ -113,8 +117,8 @@ def _create_batch(
     def publish(
         self,
         message: gapic_types.PubsubMessage,
-        retry: "retry.Retry" = gapic_v1.method.DEFAULT,
-        timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
+        retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
+        timeout: _TimeoutType = gapic_v1.method.DEFAULT,
     ) -> "futures.Future":
         """ Batch message into existing or new batch.
 

@@ -17,7 +17,7 @@
 import os
 import pkg_resources
 import typing
-from typing import Any, Callable, Optional, Sequence, Union
+from typing import cast, Any, Callable, Optional, Sequence, Union
 
 from google.auth.credentials import AnonymousCredentials  # type: ignore
 from google.oauth2 import service_account  # type: ignore
@@ -30,6 +30,7 @@
 
 if typing.TYPE_CHECKING:  # pragma: NO COVER
     from google.cloud.pubsub_v1 import subscriber
+    from google.pubsub_v1.services.subscriber.transports.grpc import SubscriberGrpcTransport
 
 
 try:
@@ -263,7 +264,8 @@ def close(self) -> None:
 
         This method is idempotent.
         """
-        self.api._transport.grpc_channel.close()
+        transport = cast("SubscriberGrpcTransport", self.api._transport)
+        transport.grpc_channel.close()
         self._closed = True
 
     def __enter__(self) -> "Client":