From a7273b9e4763b4d4dc7e56a9c27d72ce881deaf2 Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Wed, 17 Nov 2021 15:29:59 -0500 Subject: [PATCH] chore: fix lint / coverage issues --- google/cloud/firestore_v1/__init__.py | 5 +--- .../cloud/firestore_v1/async_transaction.py | 2 +- google/cloud/firestore_v1/base_client.py | 4 ++- google/cloud/firestore_v1/base_query.py | 25 ++++++------------- google/cloud/firestore_v1/base_transaction.py | 15 ++++++----- google/cloud/firestore_v1/bulk_writer.py | 12 +++++---- google/cloud/firestore_v1/query.py | 5 +--- noxfile.py | 6 +---- tests/unit/v1/test_base_document.py | 2 +- tests/unit/v1/test_bulk_writer.py | 6 +---- tests/unit/v1/test_rate_limiter.py | 1 - 11 files changed, 31 insertions(+), 52 deletions(-) diff --git a/google/cloud/firestore_v1/__init__.py b/google/cloud/firestore_v1/__init__.py index 57db5b83f8..0fade3f08f 100644 --- a/google/cloud/firestore_v1/__init__.py +++ b/google/cloud/firestore_v1/__init__.py @@ -18,15 +18,12 @@ """Python idiomatic client for Google Cloud Firestore.""" -from typing import Union - import pkg_resources -__version__: Union[str, None] try: __version__ = pkg_resources.get_distribution("google-cloud-firestore").version except pkg_resources.DistributionNotFound: - __version__ = None + __version__ = None # type: ignore from google.cloud.firestore_v1 import types from google.cloud.firestore_v1._helpers import GeoPoint diff --git a/google/cloud/firestore_v1/async_transaction.py b/google/cloud/firestore_v1/async_transaction.py index 712f045821..8cc5f0b4c8 100644 --- a/google/cloud/firestore_v1/async_transaction.py +++ b/google/cloud/firestore_v1/async_transaction.py @@ -192,7 +192,7 @@ async def get( [ref_or_query], transaction=self, **kwargs ) elif isinstance(ref_or_query, AsyncQuery): - return await ref_or_query.stream(transaction=self, **kwargs) # type: ignore + return await ref_or_query.stream(transaction=self, **kwargs) # type: ignore else: raise ValueError( 'Value for argument "ref_or_query" must be a AsyncDocumentReference or a AsyncQuery.' diff --git a/google/cloud/firestore_v1/base_client.py b/google/cloud/firestore_v1/base_client.py index 64e573636c..5e777f77f6 100644 --- a/google/cloud/firestore_v1/base_client.py +++ b/google/cloud/firestore_v1/base_client.py @@ -63,7 +63,9 @@ ) _ACTIVE_TXN: str = "There is already an active transaction." _INACTIVE_TXN: str = "There is no active transaction." -_CLIENT_INFO: Any = client_info.ClientInfo(client_library_version=__version__) +_CLIENT_INFO: Any = client_info.ClientInfo( + client_library_version=__version__ # type: ignore +) _FIRESTORE_EMULATOR_HOST: str = "FIRESTORE_EMULATOR_HOST" diff --git a/google/cloud/firestore_v1/base_query.py b/google/cloud/firestore_v1/base_query.py index 77707a9510..9e1dfb250d 100644 --- a/google/cloud/firestore_v1/base_query.py +++ b/google/cloud/firestore_v1/base_query.py @@ -45,6 +45,7 @@ from google.cloud.firestore_v1.types import RunQueryResponse from google.cloud.firestore_v1.order import Order + _BAD_DIR_STRING: str _BAD_OP_NAN_NULL: str _BAD_OP_STRING: str @@ -107,7 +108,8 @@ class _NotPassed: OptionalInt = Union[int, _NotPassed] OptionalBool = Union[bool, _NotPassed] CursorParamStripped = Tuple[Union[tuple, dict, list], bool] -CursorParam = Tuple[Union[tuple, dict, list, DocumentSnapshot], bool] +CursorArg = Union[DocumentSnapshot, dict, list, tuple, None] +CursorParam = Tuple[CursorArg, bool] OptionalCursorParam = Union[CursorParam, _NotPassed] @@ -474,10 +476,7 @@ def _check_snapshot(self, document_snapshot) -> None: raise ValueError("Cannot use snapshot from another collection as a cursor.") def _cursor_helper( - self, - document_fields_or_snapshot: Union[DocumentSnapshot, dict, list, tuple], - before: bool, - start: bool, + self, document_fields_or_snapshot: CursorArg, before: bool, start: bool, ) -> "BaseQuery": """Set values to be used for a ``start_at`` or ``end_at`` cursor. @@ -530,9 +529,7 @@ def _cursor_helper( return self._copy(**query_kwargs) - def start_at( - self, document_fields_or_snapshot: Union[DocumentSnapshot, dict, list, tuple] - ) -> "BaseQuery": + def start_at(self, document_fields_or_snapshot: CursorArg) -> "BaseQuery": """Start query results at a particular document value. The result set will **include** the document specified by @@ -562,9 +559,7 @@ def start_at( """ return self._cursor_helper(document_fields_or_snapshot, before=True, start=True) - def start_after( - self, document_fields_or_snapshot: Union[DocumentSnapshot, dict, list, tuple] - ) -> "BaseQuery": + def start_after(self, document_fields_or_snapshot: CursorArg) -> "BaseQuery": """Start query results after a particular document value. The result set will **exclude** the document specified by @@ -595,9 +590,7 @@ def start_after( document_fields_or_snapshot, before=False, start=True ) - def end_before( - self, document_fields_or_snapshot: Union[DocumentSnapshot, dict, list, tuple] - ) -> "BaseQuery": + def end_before(self, document_fields_or_snapshot: CursorArg) -> "BaseQuery": """End query results before a particular document value. The result set will **exclude** the document specified by @@ -628,9 +621,7 @@ def end_before( document_fields_or_snapshot, before=True, start=False ) - def end_at( - self, document_fields_or_snapshot: Union[DocumentSnapshot, dict, list, tuple] - ) -> "BaseQuery": + def end_at(self, document_fields_or_snapshot: CursorArg) -> "BaseQuery": """End query results at a particular document value. The result set will **include** the document specified by diff --git a/google/cloud/firestore_v1/base_transaction.py b/google/cloud/firestore_v1/base_transaction.py index 046ae7303f..83b2c1433f 100644 --- a/google/cloud/firestore_v1/base_transaction.py +++ b/google/cloud/firestore_v1/base_transaction.py @@ -14,7 +14,7 @@ """Helpers for applying Google Cloud Firestore changes in a transaction.""" -from typing import Any, Coroutine, List, NoReturn, Optional, Union +from typing import List, Union from google.cloud.firestore_v1 import types from google.cloud.firestore_v1.types import write @@ -41,10 +41,10 @@ class BaseTransaction(object): """Accumulate read-and-write operations to be sent in a transaction. Args: - max_attempts (Optional[int]): The maximum number of attempts for + max_attempts (int): The maximum number of attempts for the transaction (i.e. allowing retries). Defaults to :attr:`~google.cloud.firestore_v1.transaction.MAX_ATTEMPTS`. - read_only (Optional[bool]): Flag indicating if the transaction + read_only (bool): Flag indicating if the transaction should be read-only or should allow writes. Defaults to :data:`False`. """ @@ -57,7 +57,7 @@ def __init__(self, max_attempts=MAX_ATTEMPTS, read_only=False) -> None: def _options_protobuf( self, retry_id: Union[bytes, None] - ) -> Optional[types.common.TransactionOptions]: + ) -> Union[types.common.TransactionOptions, None]: """Convert the current object to protobuf. The ``retry_id`` value is used when retrying a transaction that @@ -69,7 +69,6 @@ def _options_protobuf( to be retried. Returns: - Optional[google.cloud.firestore_v1.types.TransactionOptions]: The protobuf ``TransactionOptions`` if ``read_only==True`` or if there is a transaction ID to be retried, else :data:`None`. @@ -107,7 +106,7 @@ def id(self): """Get the current transaction ID. Returns: - Optional[bytes]: The transaction ID (or :data:`None` if the + Union[bytes, None]: The transaction ID (or :data:`None` if the current transaction is not in progress). """ return self._id @@ -135,9 +134,9 @@ class _BaseTransactional(object): def __init__(self, to_wrap) -> None: self.to_wrap = to_wrap self.current_id = None - """Optional[bytes]: The current transaction ID.""" + """Union[bytes, None]: The current transaction ID.""" self.retry_id = None - """Optional[bytes]: The ID of the first attempted transaction.""" + """Union[bytes, None]: The ID of the first attempted transaction.""" def _reset(self) -> None: """Unset the transaction IDs.""" diff --git a/google/cloud/firestore_v1/bulk_writer.py b/google/cloud/firestore_v1/bulk_writer.py index 02d35f5c17..84bb44d06f 100644 --- a/google/cloud/firestore_v1/bulk_writer.py +++ b/google/cloud/firestore_v1/bulk_writer.py @@ -23,7 +23,7 @@ import functools import logging import time -from typing import Callable, Deque, Dict, List, Optional, Union, TYPE_CHECKING +from typing import Callable, Deque, Dict, List, Optional, Union from google.rpc import status_pb2 # type: ignore @@ -79,6 +79,7 @@ class AsyncBulkWriterMixin: The entrypoint to the parallelizable code path is `_send_batch()`, which is wrapped in a decorator which ensures that the `SendMode` is honored. """ + _in_flight_documents: int = 0 _total_batches_sent: int = 0 _total_write_operations: int = 0 @@ -503,7 +504,7 @@ def _schedule_ready_retries(self): def _request_send(self, batch_size: int) -> bool: # Set up this boolean to avoid repeatedly taking tokens if we're only # waiting on the `max_in_flight` limit. - have_received_tokens: bool = False + got_tokens: bool = False while True: # To avoid bottlenecks on the server, an additional limit is that no @@ -515,10 +516,9 @@ def _request_send(self, batch_size: int) -> bool: ) # Ask for tokens each pass through this loop until they are granted, # and then stop. - if not have_received_tokens: - have_received_tokens = bool(self._rate_limiter.take_tokens(batch_size)) + got_tokens = got_tokens or bool(self._rate_limiter.take_tokens(batch_size)) - if not under_threshold or not have_received_tokens: + if not under_threshold or not got_tokens: # Try again until both checks are true. # Note that this sleep is helpful to prevent the main BulkWriter # thread from spinning through this loop as fast as possible and @@ -727,6 +727,7 @@ class BulkWriterOperation: that ferries it into its next retry without getting confused with other similar writes to the same document. """ + attempts: int def add_to_batch(self, batch: BulkWriteBatch): @@ -766,6 +767,7 @@ class BaseOperationRetry: Methods on this class be moved directly to `OperationRetry` when support for Python 3.6 is dropped and `dataclasses` becomes universal. """ + operation: BulkWriterOperation run_at: datetime.datetime diff --git a/google/cloud/firestore_v1/query.py b/google/cloud/firestore_v1/query.py index 2b4e76ae6c..a3d54cf7c9 100644 --- a/google/cloud/firestore_v1/query.py +++ b/google/cloud/firestore_v1/query.py @@ -279,10 +279,7 @@ def stream( response = next(response_iterator, None) except exceptions.GoogleAPICallError as exc: if self._retry_query_after_exception(exc, retry, transaction): - if last_snapshot is not None: - new_query = self.start_after(last_snapshot) - else: - new_query = self + new_query = cast(Query, self.start_after(last_snapshot)) response_iterator, _ = new_query._get_stream_iterator( transaction, retry, timeout, ) diff --git a/noxfile.py b/noxfile.py index 5e413c4021..5b91b8d7d7 100644 --- a/noxfile.py +++ b/noxfile.py @@ -86,11 +86,7 @@ def mypy(session): """Verify type hints are mypy compatible.""" session.install("-e", ".") session.install( - "mypy", - "types-setuptools", - "types-protobuf", - "types-dataclasses", - "types-mock", + "mypy", "types-setuptools", "types-protobuf", "types-dataclasses", "types-mock", ) # Note: getenerated tests (in 'tests/unit/gapic') are not yet # mypy-safe diff --git a/tests/unit/v1/test_base_document.py b/tests/unit/v1/test_base_document.py index 55ed12ffb9..ba8b150458 100644 --- a/tests/unit/v1/test_base_document.py +++ b/tests/unit/v1/test_base_document.py @@ -16,7 +16,7 @@ import unittest import mock -from proto.datetime_helpers import DatetimeWithNanoseconds # type: ignore +from proto.datetime_helpers import DatetimeWithNanoseconds # type: ignore class TestBaseDocumentReference(unittest.TestCase): diff --git a/tests/unit/v1/test_bulk_writer.py b/tests/unit/v1/test_bulk_writer.py index de316adab8..59f2c83ffb 100644 --- a/tests/unit/v1/test_bulk_writer.py +++ b/tests/unit/v1/test_bulk_writer.py @@ -24,7 +24,6 @@ from google.cloud.firestore_v1.async_client import AsyncClient from google.cloud.firestore_v1.base_document import BaseDocumentReference from google.cloud.firestore_v1.client import Client -from google.cloud.firestore_v1.base_client import BaseClient from google.cloud.firestore_v1.bulk_batch import BulkWriteBatch from google.cloud.firestore_v1.bulk_writer import ( BulkRetry, @@ -139,9 +138,7 @@ def test_ctor_explicit(self): @staticmethod def _get_document_reference( - client, - collection_name: str = "col", - id: str = None, + client, collection_name: str = "col", id: str = None, ) -> Type: return client.collection(collection_name).document(id) @@ -152,7 +149,6 @@ def _doc_iter(self, client, num: int, ids: Optional[List[str]] = None): def _verify_bw_activity(self, bw: NoSendBulkWriter, counts: List[Tuple[int, int]]): total_batches = sum([el[1] for el in counts]) - batches_word = "batches" if total_batches != 1 else "batch" assert len(bw._responses) == total_batches expected_counts = dict(counts) diff --git a/tests/unit/v1/test_rate_limiter.py b/tests/unit/v1/test_rate_limiter.py index bfc65c2edc..b75c8f39da 100644 --- a/tests/unit/v1/test_rate_limiter.py +++ b/tests/unit/v1/test_rate_limiter.py @@ -14,7 +14,6 @@ import datetime import unittest -from typing import Optional import mock import google