Skip to content

feat: Add a little bit of typing to google.api_core.retry #453

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Sep 1, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .coveragerc
Original file line number Diff line number Diff line change
Expand Up @@ -11,3 +11,5 @@ exclude_lines =
def __repr__
# Ignore abstract methods
raise NotImplementedError
# Ignore coverage for code specific to static type checkers
TYPE_CHECKING
43 changes: 30 additions & 13 deletions google/api_core/retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,28 +54,41 @@ def check_if_exists():

"""

from __future__ import unicode_literals
from __future__ import annotations
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this library has python_requires >= 3.7 so (1) unicode_literals does nothing and (2) this is the first version which supports __future__.annotations

it might be a good idea to run https://github.com/asottile/pyupgrade on this codebase to remove the remaining python 2isms

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! I've captured this TODO in #524


import datetime
import functools
import logging
import random
import sys
import time
from typing import Any, Callable, TypeVar, TYPE_CHECKING

import requests.exceptions

from google.api_core import datetime_helpers
from google.api_core import exceptions
from google.auth import exceptions as auth_exceptions

if TYPE_CHECKING:
if sys.version_info >= (3, 10):
from typing import ParamSpec
else:
from typing_extensions import ParamSpec

_P = ParamSpec("_P")
_R = TypeVar("_R")

_LOGGER = logging.getLogger(__name__)
_DEFAULT_INITIAL_DELAY = 1.0 # seconds
_DEFAULT_MAXIMUM_DELAY = 60.0 # seconds
_DEFAULT_DELAY_MULTIPLIER = 2.0
_DEFAULT_DEADLINE = 60.0 * 2.0 # seconds


def if_exception_type(*exception_types):
def if_exception_type(
*exception_types: type[BaseException],
) -> Callable[[BaseException], bool]:
"""Creates a predicate to check if the exception is of a given type.

Args:
Expand All @@ -87,7 +100,7 @@ def if_exception_type(*exception_types):
exception is of the given type(s).
"""

def if_exception_type_predicate(exception):
def if_exception_type_predicate(exception: BaseException) -> bool:
"""Bound predicate for checking an exception type."""
return isinstance(exception, exception_types)

Expand Down Expand Up @@ -307,14 +320,14 @@ class Retry(object):

def __init__(
self,
predicate=if_transient_error,
initial=_DEFAULT_INITIAL_DELAY,
maximum=_DEFAULT_MAXIMUM_DELAY,
multiplier=_DEFAULT_DELAY_MULTIPLIER,
timeout=_DEFAULT_DEADLINE,
on_error=None,
**kwargs
):
predicate: Callable[[BaseException], bool] = if_transient_error,
initial: float = _DEFAULT_INITIAL_DELAY,
maximum: float = _DEFAULT_MAXIMUM_DELAY,
multiplier: float = _DEFAULT_DELAY_MULTIPLIER,
timeout: float = _DEFAULT_DEADLINE,
on_error: Callable[[BaseException], Any] | None = None,
**kwargs: Any,
) -> None:
self._predicate = predicate
self._initial = initial
self._multiplier = multiplier
Expand All @@ -323,7 +336,11 @@ def __init__(
self._deadline = self._timeout
self._on_error = on_error

def __call__(self, func, on_error=None):
def __call__(
self,
func: Callable[_P, _R],
on_error: Callable[[BaseException], Any] | None = None,
) -> Callable[_P, _R]:
"""Wrap a callable with retry behavior.

Args:
Expand All @@ -340,7 +357,7 @@ def __call__(self, func, on_error=None):
on_error = self._on_error

@functools.wraps(func)
def retry_wrapped_func(*args, **kwargs):
def retry_wrapped_func(*args: _P.args, **kwargs: _P.kwargs) -> _R:
"""A wrapper that calls target function with retry."""
target = functools.partial(func, *args, **kwargs)
sleep_generator = exponential_sleep_generator(
Expand Down