added parallelization support to Box.map

utilz/data.py

@@ -88,9 +88,10 @@ def fn(*args, **kwargs):
     def __repr__(self):
         return f"Box(len={len(self)}, transparent={self._transparent_box}, type={self[0].__class__.__module__}.{self[0].__class__.__name__})"
 
-    def map(self, fn, inplace=False):
+    def map(self, fn, **kwargs):
         """
-        Apply a function to each element in the box
+        Apply a function to each element in the box. Accepts all kwargs that map does,
+        including parallelization!
 
         Args:
             fn (callable): function to apply to each element
@@ -102,7 +103,8 @@ def map(self, fn, inplace=False):
             Box: new box with the results of applying `fn` to each element
         """
 
-        out = map(fn, self)
+        inplace = kwargs.pop("inplace", False)
+        out = map(fn, self, **kwargs)
         if inplace:
             self.__init__(out, transparent=self._transparent_box)
         else:

utilz/tests/test_data.py

@@ -3,6 +3,8 @@
 from utilz.data import Box
 from utilz import randdf, equal, map, seq
 import pytest
+from timeit import timeit
+from time import sleep
 
 
 def test_box():
@@ -100,3 +102,18 @@ def test_box():
     box = Box(data, transparent=False)
     box.map(lambda x: x + 1, inplace=True)
     assert equal(box.contents(), correct)
+
+
+def test_parallel_box():
+    df_data = [randdf((20, 3), groups={"condition": 2, "group": 4}) for i in range(4)]
+
+    def sleepy_mean(df):
+        sleep(1)
+        return df.mean()
+
+    box = Box(df_data)
+
+    single = timeit(lambda: box.map(sleepy_mean), number=1)
+    par = timeit(lambda: box.map(sleepy_mean, n_jobs=2), number=1)
+
+    assert single > par