feat: add shuffle option in read_files

Author: flavioschneider

pipd/functions/read_files.py

@@ -1,6 +1,7 @@
 import asyncio
 import glob
 import os
+import random
 from typing import Iterable, Iterator, Optional, Sequence, TypeVar
 
 from pipd import Function, Pipe
@@ -29,10 +30,14 @@ def watchdir(
 
 class ReadFiles(Function):
     def __init__(
-        self, cache_filepath: Optional[str] = None, watch: bool = False
+        self,
+        cache_filepath: Optional[str] = None,
+        watch: bool = False,
+        shuffle: bool = False,
     ) -> None:
         self.cache_filepath = cache_filepath
         self.watch = watch
+        self.shuffle = shuffle
 
     def __call__(self, items: Iterable[str]) -> Iterator[str]:
         for filepath in items:
@@ -46,6 +51,9 @@ def __call__(self, items: Iterable[str]) -> Iterator[str]:
             else:
                 files = glob.glob(filepath)
 
+            if self.shuffle:
+                random.shuffle(files)
+
             for file in files:
                 yield file
             if self.watch:

pipd/tests/test_functions.py

@@ -158,3 +158,14 @@ def test_read_files():
         pipe = Pipe([f"{d}/*.txt"]).read_files(cache_filepath=f"{d}/cache.txt")
         assert list(pipe) == [f1.name, f2.name]
         os.remove(f"{d}/cache.txt")
+
+    # Test shuffle
+    with tempfile.TemporaryDirectory() as d:
+        with open(os.path.join(d, "f1.txt"), "w") as f1:
+            f1.write("a")
+        with open(os.path.join(d, "f2.txt"), "w") as f2:
+            f2.write("b")
+        pipe = Pipe([f"{d}/*.txt"]).read_files(shuffle=True)
+        assert sorted(list(pipe)) == sorted([f1.name, f2.name])
+        os.remove(f1.name)
+        os.remove(f2.name)

setup.py

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