feat: add partial option to disable partial batches

elevenlabs · Jun 6, 2023 · 8537fe1 · 8537fe1
1 parent 7e148ab
commit 8537fe1
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/pipd/functions/batch.py b/pipd/functions/batch.py
@@ -6,8 +6,9 @@
 
 
 class Batch(Function):
-    def __init__(self, size: int) -> None:
+    def __init__(self, size: int, partial: bool = True) -> None:
         self.size = size
+        self.partial = partial
 
     def __call__(self, items: Iterable[T]) -> Iterator[List[T]]:
         batch = []
@@ -16,7 +17,7 @@ def __call__(self, items: Iterable[T]) -> Iterator[List[T]]:
             if len(batch) == self.size:
                 yield batch
                 batch = []
-        if batch:
+        if batch and self.partial:
             yield batch
 
 

diff --git a/pipd/tests/test_functions.py b/pipd/tests/test_functions.py
@@ -37,6 +37,9 @@ def test_batch():
     pipe = Pipe(range(5)).batch(2)
     assert list(pipe) == [[0, 1], [2, 3], [4]]
 
+    pipe = Pipe(range(5)).batch(2, partial=False)
+    assert list(pipe) == [[0, 1], [2, 3]]
+
 
 def test_unbatch():
     pipe = Pipe([[0, 1], [2, 3], [4]]).unbatch()

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
     name="pipd",
     packages=find_packages(exclude=[]),
-    version="0.1.2",
+    version="0.1.3",
     description="Utility functions for python data pipelines.",
     long_description_content_type="text/markdown",
     author="ElevenLabs",