Apply extra_kwargs in batches wait_for_results

ntropy_sdk/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "5.1.1"
+__version__ = "5.1.2"
 
 from typing import TYPE_CHECKING, Optional
 

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:

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 5.1.1
+current_version = 5.1.2
 commit = True
 tag = True
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(rc(?P<pre>\d+))?

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,
 )