diff --git a/ntropy_sdk/__init__.py b/ntropy_sdk/__init__.py index 1fc25b3..acb6186 100644 --- a/ntropy_sdk/__init__.py +++ b/ntropy_sdk/__init__.py @@ -1,4 +1,4 @@ -__version__ = "5.1.1" +__version__ = "5.1.2" from typing import TYPE_CHECKING, Optional diff --git a/ntropy_sdk/batches.py b/ntropy_sdk/batches.py index ac73ccd..5f0fb46 100644 --- a/ntropy_sdk/batches.py +++ b/ntropy_sdk/batches.py @@ -174,11 +174,12 @@ def _wait( poll_interval: int, timeout: int, stop_fn: Callable[[Batch], bool], + extra_kwargs: "ExtraKwargs", ) -> Batch: start_time = time.monotonic() batch = None while time.monotonic() - start_time < timeout: - batch = self.get(id=id) + batch = self.get(id=id, **extra_kwargs) if stop_fn(batch): break time.sleep(poll_interval) @@ -191,6 +192,7 @@ def _wait_with_progress( poll_interval: int, timeout: int, stop_fn: Callable[[Batch], bool], + extra_kwargs: "ExtraKwargs", ) -> Batch: from tqdm.auto import tqdm @@ -199,7 +201,7 @@ def _wait_with_progress( total_set = False with tqdm() as p: while time.monotonic() - start_time < timeout: - batch = self.get(id=id) + batch = self.get(id=id, **extra_kwargs) if not total_set: p.total = batch.total p.desc = batch.status @@ -230,11 +232,19 @@ def stop_fn(b: Batch): if with_progress: batch = self._wait_with_progress( - id=id, poll_interval=poll_interval, timeout=timeout, stop_fn=stop_fn + id=id, + poll_interval=poll_interval, + timeout=timeout, + stop_fn=stop_fn, + extra_kwargs=extra_kwargs, ) else: batch = self._wait( - id=id, poll_interval=poll_interval, timeout=timeout, stop_fn=stop_fn + id=id, + poll_interval=poll_interval, + timeout=timeout, + stop_fn=stop_fn, + extra_kwargs=extra_kwargs, ) if batch and batch.status not in finish_statuses: diff --git a/setup.cfg b/setup.cfg index c539062..c05cd1d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 5.1.1 +current_version = 5.1.2 commit = True tag = True parse = (?P\d+)\.(?P\d+)\.(?P\d+)(rc(?P
\d+))?
diff --git a/setup.py b/setup.py
index 3eb5db3..bc83dad 100644
--- a/setup.py
+++ b/setup.py
@@ -55,6 +55,6 @@
     test_suite="tests",
     tests_require=test_requirements,
     url="https://github.com/ntropy-network/ntropy-sdk",
-    version="5.1.1",
+    version="5.1.2",
     zip_safe=False,
 )