Skip to content

Commit 8f32dd4

Browse files
authored
process: make sure type annoations pass with mypy (#542)
* Add mypy to nox sessions * Mark google/cloud package as type-checked * Ignore types for dependencies lacking type info * Fix type annotations in publish flow controller * Fix type hints in thread-based Batch * Fix type annotations in sequencers * Fix type hints in dispatcher * Fix type hints in publisher client * Fix misc type errors in various modules * Fix type hints in leaser * Fix type annotations in streaming pull manager * Fix gapic timeout hint in older api-core versions google-api-core versions prior to v2.2.2 lack the definition of _MethodDefault, thus a workaround is needed for that. * Remove py.typed marker file The autogenerated code does not pass mypy type checks yet, thus we should not advertise the package as type-checked. * Replace typing.cast with is not None assertions * Replace batched_commands dict with separate lists * Rename variable to avoid false type warnings * Get rid of type cast by using a new variable * Just ignore the line where type checkers disagree * Remove unused imports * Replace type casts with is not None assertions * Cover missing dispatcher case after refactoring
1 parent c1b424e commit 8f32dd4

File tree

20 files changed

+240
-130
lines changed

20 files changed

+240
-130
lines changed

google/cloud/__init__.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,13 @@
1414
# See the License for the specific language governing permissions and
1515
# limitations under the License.
1616

17+
from typing import List
18+
1719
try:
1820
import pkg_resources
1921

2022
pkg_resources.declare_namespace(__name__)
2123
except ImportError:
2224
import pkgutil
2325

24-
__path__ = pkgutil.extend_path(__path__, __name__)
26+
__path__: List[str] = pkgutil.extend_path(__path__, __name__) # type: ignore

google/cloud/pubsub_v1/futures.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
from __future__ import absolute_import
1616

1717
import concurrent.futures
18-
from typing import Any, NoReturn
18+
from typing import Any, NoReturn, Optional
1919

2020
import google.api_core.future
2121

@@ -47,7 +47,7 @@ def set_result(self, result: Any):
4747
"""
4848
return super().set_result(result=result)
4949

50-
def set_exception(self, exception: Exception):
50+
def set_exception(self, exception: Optional[BaseException]):
5151
"""Set the result of the future as being the given exception.
5252
5353
Do not use this method, it should only be used internally by the library and its

google/cloud/pubsub_v1/publisher/_batch/base.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ def __len__(self):
5858

5959
@staticmethod
6060
@abc.abstractmethod
61-
def make_lock() -> None: # pragma: NO COVER
61+
def make_lock(): # pragma: NO COVER
6262
"""Return a lock in the chosen concurrency model.
6363
6464
Returns:

google/cloud/pubsub_v1/publisher/_batch/thread.py

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
import threading
1919
import time
2020
import typing
21-
from typing import Any, Callable, Optional, Sequence
21+
from typing import Any, Callable, List, Optional, Sequence
2222

2323
import google.api_core.exceptions
2424
from google.api_core import gapic_v1
@@ -28,11 +28,10 @@
2828
from google.pubsub_v1 import types as gapic_types
2929

3030
if typing.TYPE_CHECKING: # pragma: NO COVER
31-
from google import api_core
3231
from google.cloud import pubsub_v1
3332
from google.cloud.pubsub_v1 import types
34-
from google.cloud.pubsub_v1 import PublisherClient
35-
33+
from google.cloud.pubsub_v1.publisher import Client as PublisherClient
34+
from google.pubsub_v1.services.publisher.client import OptionalRetry
3635

3736
_LOGGER = logging.getLogger(__name__)
3837
_CAN_COMMIT = (base.BatchStatus.ACCEPTING_MESSAGES, base.BatchStatus.STARTING)
@@ -93,8 +92,8 @@ def __init__(
9392
settings: "types.BatchSettings",
9493
batch_done_callback: Callable[[bool], Any] = None,
9594
commit_when_full: bool = True,
96-
commit_retry: "api_core.retry.Retry" = gapic_v1.method.DEFAULT,
97-
commit_timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
95+
commit_retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
96+
commit_timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
9897
):
9998
self._client = client
10099
self._topic = topic
@@ -108,8 +107,8 @@ def __init__(
108107
# _futures list should remain unchanged after batch
109108
# status changed from ACCEPTING_MESSAGES to any other
110109
# in order to avoid race conditions
111-
self._futures = []
112-
self._messages = []
110+
self._futures: List[futures.Future] = []
111+
self._messages: List[gapic_types.PubsubMessage] = []
113112
self._status = base.BatchStatus.ACCEPTING_MESSAGES
114113

115114
# The initial size is not zero, we need to account for the size overhead
@@ -368,7 +367,7 @@ def publish(
368367
), "Publish after stop() or publish error."
369368

370369
if self.status != base.BatchStatus.ACCEPTING_MESSAGES:
371-
return
370+
return None
372371

373372
size_increase = gapic_types.PublishRequest(
374373
messages=[message]

google/cloud/pubsub_v1/publisher/_sequencer/base.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,15 +22,14 @@
2222

2323
if typing.TYPE_CHECKING: # pragma: NO COVER
2424
from concurrent import futures
25-
from google.api_core import retry
25+
from google.pubsub_v1.services.publisher.client import OptionalRetry
2626

2727

2828
class Sequencer(metaclass=abc.ABCMeta):
2929
"""The base class for sequencers for Pub/Sub publishing. A sequencer
3030
sequences messages to be published.
3131
"""
3232

33-
@staticmethod
3433
@abc.abstractmethod
3534
def is_finished(self) -> bool: # pragma: NO COVER
3635
""" Whether the sequencer is finished and should be cleaned up.
@@ -40,7 +39,6 @@ def is_finished(self) -> bool: # pragma: NO COVER
4039
"""
4140
raise NotImplementedError
4241

43-
@staticmethod
4442
@abc.abstractmethod
4543
def unpause(self) -> None: # pragma: NO COVER
4644
""" Unpauses this sequencer.
@@ -51,12 +49,11 @@ def unpause(self) -> None: # pragma: NO COVER
5149
"""
5250
raise NotImplementedError
5351

54-
@staticmethod
5552
@abc.abstractmethod
5653
def publish(
5754
self,
5855
message: gapic_types.PubsubMessage,
59-
retry: "retry.Retry" = None,
56+
retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
6057
timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
6158
) -> "futures.Future": # pragma: NO COVER
6259
""" Publish message for this ordering key.

google/cloud/pubsub_v1/publisher/_sequencer/ordered_sequencer.py

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,21 +14,22 @@
1414

1515
import enum
1616
import collections
17-
import concurrent.futures as futures
1817
import threading
1918
import typing
20-
from typing import Iterable, Sequence
19+
from typing import Deque, Iterable, Sequence
2120

2221
from google.api_core import gapic_v1
22+
from google.cloud.pubsub_v1.publisher import futures
2323
from google.cloud.pubsub_v1.publisher import exceptions
2424
from google.cloud.pubsub_v1.publisher._sequencer import base as sequencer_base
2525
from google.cloud.pubsub_v1.publisher._batch import base as batch_base
2626
from google.pubsub_v1 import types as gapic_types
2727

2828
if typing.TYPE_CHECKING: # pragma: NO COVER
29-
from google.api_core import retry
30-
from google.cloud.pubsub_v1 import PublisherClient
29+
from google.cloud.pubsub_v1 import types
3130
from google.cloud.pubsub_v1.publisher import _batch
31+
from google.cloud.pubsub_v1.publisher.client import Client as PublisherClient
32+
from google.pubsub_v1.services.publisher.client import OptionalRetry
3233

3334

3435
class _OrderedSequencerStatus(str, enum.Enum):
@@ -101,7 +102,7 @@ def __init__(self, client: "PublisherClient", topic: str, ordering_key: str):
101102
# Batches ordered from first (head/left) to last (right/tail).
102103
# Invariant: always has at least one batch after the first publish,
103104
# unless paused or stopped.
104-
self._ordered_batches = collections.deque()
105+
self._ordered_batches: Deque["_batch.thread.Batch"] = collections.deque()
105106
# See _OrderedSequencerStatus for valid state transitions.
106107
self._state = _OrderedSequencerStatus.ACCEPTING_MESSAGES
107108

@@ -237,8 +238,8 @@ def unpause(self) -> None:
237238

238239
def _create_batch(
239240
self,
240-
commit_retry: "retry.Retry" = gapic_v1.method.DEFAULT,
241-
commit_timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
241+
commit_retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
242+
commit_timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
242243
) -> "_batch.thread.Batch":
243244
""" Create a new batch using the client's batch class and other stored
244245
settings.
@@ -262,8 +263,8 @@ def _create_batch(
262263
def publish(
263264
self,
264265
message: gapic_types.PubsubMessage,
265-
retry: "retry.Retry" = gapic_v1.method.DEFAULT,
266-
timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
266+
retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
267+
timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
267268
) -> futures.Future:
268269
""" Publish message for this ordering key.
269270
@@ -289,12 +290,12 @@ def publish(
289290
"""
290291
with self._state_lock:
291292
if self._state == _OrderedSequencerStatus.PAUSED:
292-
future = futures.Future()
293+
errored_future = futures.Future()
293294
exception = exceptions.PublishToPausedOrderingKeyException(
294295
self._ordering_key
295296
)
296-
future.set_exception(exception)
297-
return future
297+
errored_future.set_exception(exception)
298+
return errored_future
298299

299300
# If waiting to be cleaned-up, convert to accepting messages to
300301
# prevent this sequencer from being cleaned-up only to have another

google/cloud/pubsub_v1/publisher/_sequencer/unordered_sequencer.py

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,17 +13,20 @@
1313
# limitations under the License.
1414

1515
import typing
16+
from typing import Optional
1617

1718
from google.api_core import gapic_v1
1819

1920
from google.cloud.pubsub_v1.publisher._sequencer import base
2021
from google.pubsub_v1 import types as gapic_types
2122

2223
if typing.TYPE_CHECKING: # pragma: NO COVER
23-
from concurrent import futures
24-
from google.api_core import retry
25-
from google.cloud.pubsub_v1 import PublisherClient
2624
from google.cloud.pubsub_v1.publisher import _batch
25+
from google.cloud.pubsub_v1.publisher import futures
26+
from google.cloud.pubsub_v1.publisher.client import Client as PublisherClient
27+
from google.pubsub_v1.services.publisher.client import OptionalRetry
28+
29+
from google.cloud.pubsub_v1 import types
2730

2831

2932
class UnorderedSequencer(base.Sequencer):
@@ -35,7 +38,7 @@ class UnorderedSequencer(base.Sequencer):
3538
def __init__(self, client: "PublisherClient", topic: str):
3639
self._client = client
3740
self._topic = topic
38-
self._current_batch = None
41+
self._current_batch: Optional["_batch.thread.Batch"] = None
3942
self._stopped = False
4043

4144
def is_finished(self) -> bool:
@@ -88,8 +91,8 @@ def unpause(self) -> typing.NoReturn:
8891

8992
def _create_batch(
9093
self,
91-
commit_retry: "retry.Retry" = gapic_v1.method.DEFAULT,
92-
commit_timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
94+
commit_retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
95+
commit_timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
9396
) -> "_batch.thread.Batch":
9497
""" Create a new batch using the client's batch class and other stored
9598
settings.
@@ -113,8 +116,8 @@ def _create_batch(
113116
def publish(
114117
self,
115118
message: gapic_types.PubsubMessage,
116-
retry: "retry.Retry" = gapic_v1.method.DEFAULT,
117-
timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
119+
retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
120+
timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
118121
) -> "futures.Future":
119122
""" Batch message into existing or new batch.
120123

google/cloud/pubsub_v1/publisher/client.py

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,11 @@
2121
import threading
2222
import time
2323
import typing
24-
from typing import Any, Sequence, Type, Union
24+
from typing import Any, Dict, Optional, Sequence, Tuple, Type, Union
2525

2626
from google.api_core import gapic_v1
27-
from google.auth.credentials import AnonymousCredentials
28-
from google.oauth2 import service_account
27+
from google.auth.credentials import AnonymousCredentials # type: ignore
28+
from google.oauth2 import service_account # type: ignore
2929

3030
from google.cloud.pubsub_v1 import _gapic
3131
from google.cloud.pubsub_v1 import types
@@ -46,10 +46,9 @@
4646
__version__ = "0.0"
4747

4848
if typing.TYPE_CHECKING: # pragma: NO COVER
49-
from google import api_core
5049
from google.cloud import pubsub_v1
51-
from google.cloud.pubsub_v1.publisher._sequencer.base import Sequencer
5250
from google.cloud.pubsub_v1.publisher import _batch
51+
from google.pubsub_v1.services.publisher.client import OptionalRetry
5352

5453

5554
_LOGGER = logging.getLogger(__name__)
@@ -62,6 +61,10 @@
6261

6362
_raw_proto_pubbsub_message = gapic_types.PubsubMessage.pb()
6463

64+
SequencerType = Union[
65+
ordered_sequencer.OrderedSequencer, unordered_sequencer.UnorderedSequencer
66+
]
67+
6568

6669
@_gapic.add_methods(publisher_client.PublisherClient, denylist=_DENYLISTED_METHODS)
6770
class Client(object):
@@ -152,10 +155,10 @@ def __init__(
152155
# messages. One batch exists for each topic.
153156
self._batch_lock = self._batch_class.make_lock()
154157
# (topic, ordering_key) => sequencers object
155-
self._sequencers = {}
158+
self._sequencers: Dict[Tuple[str, str], SequencerType] = {}
156159
self._is_stopped = False
157160
# Thread created to commit all sequencers after a timeout.
158-
self._commit_thread = None
161+
self._commit_thread: Optional[threading.Thread] = None
159162

160163
# The object controlling the message publishing flow
161164
self._flow_controller = FlowController(self.publisher_options.flow_control)
@@ -196,7 +199,7 @@ def target(self) -> str:
196199
"""
197200
return self._target
198201

199-
def _get_or_create_sequencer(self, topic: str, ordering_key: str) -> "Sequencer":
202+
def _get_or_create_sequencer(self, topic: str, ordering_key: str) -> SequencerType:
200203
""" Get an existing sequencer or create a new one given the (topic,
201204
ordering_key) pair.
202205
"""
@@ -254,8 +257,8 @@ def publish(
254257
topic: str,
255258
data: bytes,
256259
ordering_key: str = "",
257-
retry: "api_core.retry.Retry" = gapic_v1.method.DEFAULT,
258-
timeout: gapic_types.TimeoutType = gapic_v1.method.DEFAULT,
260+
retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
261+
timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
259262
**attrs: Union[bytes, str],
260263
) -> "pubsub_v1.publisher.futures.Future":
261264
"""Publish a single message.
@@ -380,8 +383,10 @@ def on_publish_done(future):
380383
if retry is gapic_v1.method.DEFAULT:
381384
# use the default retry for the publish GRPC method as a base
382385
transport = self.api._transport
383-
retry = transport._wrapped_methods[transport.publish]._retry
384-
retry = retry.with_deadline(2.0 ** 32)
386+
base_retry = transport._wrapped_methods[transport.publish]._retry
387+
retry = base_retry.with_deadline(2.0 ** 32)
388+
else:
389+
retry = retry.with_deadline(2.0 ** 32)
385390

386391
# Delegate the publishing to the sequencer.
387392
sequencer = self._get_or_create_sequencer(topic, ordering_key)
@@ -490,7 +495,7 @@ def _set_batch_class(self, batch_class: Type) -> None:
490495

491496
# Used only for testing.
492497
def _set_sequencer(
493-
self, topic: str, sequencer: "Sequencer", ordering_key: str = ""
498+
self, topic: str, sequencer: SequencerType, ordering_key: str = ""
494499
) -> None:
495500
sequencer_key = (topic, ordering_key)
496501
self._sequencers[sequencer_key] = sequencer

0 commit comments

Comments
 (0)