fix: Retry constructors methods support None (#592)

3 files changed, 77 insertions, 49 deletions

diff --git a/google/api_core/retry/retry_base.py b/google/api_core/retry/retry_base.py
index efd6d8f..1606e0f 100644
--- a/google/api_core/retry/retry_base.py
+++ b/google/api_core/retry/retry_base.py
@@ -25,7 +25,7 @@ import random
 import time
 
 from enum import Enum
-from typing import Any, Callable, TYPE_CHECKING
+from typing import Any, Callable, Optional, TYPE_CHECKING
 
 import requests.exceptions
 
@@ -238,8 +238,8 @@ class _BaseRetry(object):
         initial: float = _DEFAULT_INITIAL_DELAY,
         maximum: float = _DEFAULT_MAXIMUM_DELAY,
         multiplier: float = _DEFAULT_DELAY_MULTIPLIER,
-        timeout: float = _DEFAULT_DEADLINE,
-        on_error: Callable[[Exception], Any] | None = None,
+        timeout: Optional[float] = _DEFAULT_DEADLINE,
+        on_error: Optional[Callable[[Exception], Any]] = None,
         **kwargs: Any,
     ) -> None:
         self._predicate = predicate
@@ -265,24 +265,6 @@ class _BaseRetry(object):
     def timeout(self) -> float | None:
         return self._timeout
 
-    def _replace(
-        self,
-        predicate: Callable[[Exception], bool] | None = None,
-        initial: float | None = None,
-        maximum: float | None = None,
-        multiplier: float | None = None,
-        timeout: float | None = None,
-        on_error: Callable[[Exception], Any] | None = None,
-    ) -> Self:
-        return type(self)(
-            predicate=predicate or self._predicate,
-            initial=initial or self._initial,
-            maximum=maximum or self._maximum,
-            multiplier=multiplier or self._multiplier,
-            timeout=timeout or self._timeout,
-            on_error=on_error or self._on_error,
-        )
-
     def with_deadline(self, deadline: float | None) -> Self:
         """Return a copy of this retry with the given timeout.
 
@@ -290,23 +272,32 @@ class _BaseRetry(object):
         documentation for details.
 
         Args:
-            deadline (float): How long to keep retrying, in seconds.
+            deadline (float|None): How long to keep retrying, in seconds. If None,
+                no timeout is enforced.
 
         Returns:
             Retry: A new retry instance with the given timeout.
         """
-        return self._replace(timeout=deadline)
+        return self.with_timeout(deadline)
 
-    def with_timeout(self, timeout: float) -> Self:
+    def with_timeout(self, timeout: float | None) -> Self:
         """Return a copy of this retry with the given timeout.
 
         Args:
-            timeout (float): How long to keep retrying, in seconds.
+            timeout (float): How long to keep retrying, in seconds. If None,
+                no timeout will be enforced.
 
         Returns:
             Retry: A new retry instance with the given timeout.
         """
-        return self._replace(timeout=timeout)
+        return type(self)(
+            predicate=self._predicate,
+            initial=self._initial,
+            maximum=self._maximum,
+            multiplier=self._multiplier,
+            timeout=timeout,
+            on_error=self._on_error,
+        )
 
     def with_predicate(self, predicate: Callable[[Exception], bool]) -> Self:
         """Return a copy of this retry with the given predicate.
@@ -318,26 +309,42 @@ class _BaseRetry(object):
         Returns:
             Retry: A new retry instance with the given predicate.
         """
-        return self._replace(predicate=predicate)
+        return type(self)(
+            predicate=predicate,
+            initial=self._initial,
+            maximum=self._maximum,
+            multiplier=self._multiplier,
+            timeout=self._timeout,
+            on_error=self._on_error,
+        )
 
     def with_delay(
         self,
-        initial: float | None = None,
-        maximum: float | None = None,
-        multiplier: float | None = None,
+        initial: Optional[float] = None,
+        maximum: Optional[float] = None,
+        multiplier: Optional[float] = None,
     ) -> Self:
         """Return a copy of this retry with the given delay options.
 
         Args:
             initial (float): The minimum amount of time to delay (in seconds). This must
-                be greater than 0.
-            maximum (float): The maximum amount of time to delay (in seconds).
-            multiplier (float): The multiplier applied to the delay.
+                be greater than 0. If None, the current value is used.
+            maximum (float): The maximum amount of time to delay (in seconds). If None, the
+                current value is used.
+            multiplier (float): The multiplier applied to the delay. If None, the current
+                value is used.
 
         Returns:
-            Retry: A new retry instance with the given predicate.
+            Retry: A new retry instance with the given delay options.
         """
-        return self._replace(initial=initial, maximum=maximum, multiplier=multiplier)
+        return type(self)(
+            predicate=self._predicate,
+            initial=initial if initial is not None else self._initial,
+            maximum=maximum if maximum is not None else self._maximum,
+            multiplier=multiplier if multiplier is not None else self._multiplier,
+            timeout=self._timeout,
+            on_error=self._on_error,
+        )
 
     def __str__(self) -> str:
         return (
diff --git a/google/api_core/retry/retry_unary.py b/google/api_core/retry/retry_unary.py
index ae59d51..ab1b403 100644
--- a/google/api_core/retry/retry_unary.py
+++ b/google/api_core/retry/retry_unary.py
@@ -251,7 +251,7 @@ class Retry(_BaseRetry):
             must be greater than 0.
         maximum (float): The maximum amount of time to delay in seconds.
         multiplier (float): The multiplier applied to the delay.
-        timeout (float): How long to keep retrying, in seconds.
+        timeout (Optional[float]): How long to keep retrying, in seconds.
             Note: timeout is only checked before initiating a retry, so the target may
             run past the timeout value as long as it is healthy.
         on_error (Callable[Exception]): A function to call while processing
diff --git a/tests/unit/retry/test_retry_base.py b/tests/unit/retry/test_retry_base.py
index fa55d93..a0c6776 100644
--- a/tests/unit/retry/test_retry_base.py
+++ b/tests/unit/retry/test_retry_base.py
@@ -138,7 +138,8 @@ class Test_BaseRetry(object):
         assert retry_._on_error is _some_function
 
     @pytest.mark.parametrize("use_deadline", [True, False])
-    def test_with_timeout(self, use_deadline):
+    @pytest.mark.parametrize("value", [None, 0, 1, 4, 42, 5.5])
+    def test_with_timeout(self, use_deadline, value):
         retry_ = self._make_one(
             predicate=mock.sentinel.predicate,
             initial=1,
@@ -148,11 +149,17 @@ class Test_BaseRetry(object):
             on_error=mock.sentinel.on_error,
         )
         new_retry = (
-            retry_.with_timeout(42) if not use_deadline else retry_.with_deadline(42)
+            retry_.with_timeout(value)
+            if not use_deadline
+            else retry_.with_deadline(value)
         )
         assert retry_ is not new_retry
-        assert new_retry._timeout == 42
-        assert new_retry.timeout == 42 if not use_deadline else new_retry.deadline == 42
+        assert new_retry._timeout == value
+        assert (
+            new_retry.timeout == value
+            if not use_deadline
+            else new_retry.deadline == value
+        )
 
         # the rest of the attributes should remain the same
         assert new_retry._predicate is retry_._predicate
@@ -196,20 +203,34 @@ class Test_BaseRetry(object):
         assert new_retry._maximum == retry_._maximum
         assert new_retry._multiplier == retry_._multiplier
 
-    def test_with_delay(self):
+    @pytest.mark.parametrize(
+        "originals,updated,expected",
+        [
+            [(1, 2, 3), (4, 5, 6), (4, 5, 6)],
+            [(1, 2, 3), (0, 0, 0), (0, 0, 0)],
+            [(1, 2, 3), (None, None, None), (1, 2, 3)],
+            [(0, 0, 0), (None, None, None), (0, 0, 0)],
+            [(1, 2, 3), (None, 0.5, None), (1, 0.5, 3)],
+            [(1, 2, 3), (None, 0.5, 4), (1, 0.5, 4)],
+            [(1, 2, 3), (9, None, None), (9, 2, 3)],
+        ],
+    )
+    def test_with_delay(self, originals, updated, expected):
         retry_ = self._make_one(
             predicate=mock.sentinel.predicate,
-            initial=1,
-            maximum=2,
-            multiplier=3,
-            timeout=4,
+            initial=originals[0],
+            maximum=originals[1],
+            multiplier=originals[2],
+            timeout=14,
             on_error=mock.sentinel.on_error,
         )
-        new_retry = retry_.with_delay(initial=5, maximum=6, multiplier=7)
+        new_retry = retry_.with_delay(
+            initial=updated[0], maximum=updated[1], multiplier=updated[2]
+        )
         assert retry_ is not new_retry
-        assert new_retry._initial == 5
-        assert new_retry._maximum == 6
-        assert new_retry._multiplier == 7
+        assert new_retry._initial == expected[0]
+        assert new_retry._maximum == expected[1]
+        assert new_retry._multiplier == expected[2]
 
         # the rest of the attributes should remain the same
         assert new_retry._timeout == retry_._timeout