wip

Author: jooola

hcloud/_utils.py

@@ -3,6 +3,8 @@
 from collections.abc import Iterable, Iterator
 from itertools import islice
 from typing import TypeVar
+from functools import wraps
+from concurrent.futures import ThreadPoolExecutor
 
 T = TypeVar("T")
 
@@ -17,3 +19,18 @@ def batched(iterable: Iterable[T], size: int) -> Iterator[tuple[T, ...]]:
         if not batch:
             break
         yield batch
+
+
+def with_timeout(seconds: float | None):
+    def decorator(f):
+        @wraps(f)
+        def wrapper(*args, **kwargs):
+            if seconds:
+                with ThreadPoolExecutor(max_workers=1) as executor:
+                    future = executor.submit(f, *args, **kwargs)
+                    return future.result(timeout=seconds)
+            else:
+                return f(*args, **kwargs)
+
+        return wrapper
+    return decorator

hcloud/actions/client.py

@@ -2,9 +2,11 @@
 
 import time
 import warnings
+from concurrent.futures import ThreadPoolExecutor
+from threading import Event
 from typing import TYPE_CHECKING, Any, Callable, NamedTuple
 
-from .._utils import batched
+from .._utils import batched, with_timeout
 from ..core import BoundModelBase, ClientEntityBase, Meta
 from .domain import Action, ActionFailedException, ActionTimeoutException
 
@@ -168,6 +170,7 @@ def _get_list_by_ids(self, ids: list[int]) -> list[BoundAction]:
 
     def wait_for_function(
         self,
+        stop: Event,
         handle_update: Callable[[BoundAction], None],
         actions: list[Action | BoundAction],
     ) -> list[BoundAction]:
@@ -190,7 +193,9 @@ def wait_for_function(
         retries = 0
         while len(running_ids):
             # pylint: disable=protected-access
-            time.sleep(self._client._poll_interval_func(retries))
+            if stop.wait(self._client._poll_interval_func(retries)):
+                raise ActionTimeoutException()
+
             retries += 1
 
             updates = self._get_list_by_ids(running_ids)
@@ -207,6 +212,7 @@ def wait_for_function(
     def wait_for(
         self,
         actions: list[Action | BoundAction],
+        timeout: float | None = None,
     ) -> list[BoundAction]:
         """
         Waits until all Actions succeed by polling the API at the interval defined by
@@ -217,14 +223,18 @@ def wait_for(
 
         :param actions: List of Actions to wait for.
         :raises: ActionFailedException when an Action failed.
+        :raises: TimeoutError when the Actions did not succeed before timeout.
         :return: List of succeeded Actions.
         """
 
         def handle_update(update: BoundAction) -> None:
             if update.status == Action.STATUS_ERROR:
                 raise ActionFailedException(action=update)
 
-        return self.wait_for_function(handle_update, actions)
+        def run() -> list[BoundAction]:
+            return self.wait_for_function(handle_update, actions)
+
+        return with_timeout(timeout)(run)()
 
     def get_list(
         self,