diff --git a/dascore/core/spool.py b/dascore/core/spool.py index ac0c02895..764c0c30b 100644 --- a/dascore/core/spool.py +++ b/dascore/core/spool.py @@ -332,11 +332,11 @@ def split( ---------- size The number of patches desired in each output spool. The last - spool may have fewer patches. + spool may have fewer patches. Must be greater than zero. count The number of spools to include. If count is greater than the length of the spool then the output will be smaller than - count, with one patch per spool. + count, with one patch per spool. Must be greater than zero. Examples -------- @@ -583,12 +583,18 @@ def split( if not ((count is not None) ^ (size is not None)): msg = "Spool.split requires either spool_count or spool_size." raise ParameterError(msg) + value = count if count is not None else size + assert value is not None # the check above sets exactly one of them + # A step of zero or less never advances start, so the loop below + # would yield forever. + if value <= 0: + msg = f"Spool.split requires a positive size or count, got {value}." + raise ParameterError(msg) start = 0 if count is not None: - step = int(np.ceil(len(self) / count)) + step = int(np.ceil(len(self) / value)) else: - assert size is not None # the check above sets exactly one of them - step = int(np.ceil(size)) # tolerate a non-integral size + step = int(np.ceil(value)) # tolerate a non-integral size while start < len(self): yield self[start : start + step] start += step diff --git a/dascore/utils/misc.py b/dascore/utils/misc.py index 47a147e5d..5e945c78d 100644 --- a/dascore/utils/misc.py +++ b/dascore/utils/misc.py @@ -771,8 +771,9 @@ def _spool_map(spool, func, size=None, client=None, progress=True, **kwargs): # Now things get interesting. We need to split the spool here # so that patches don't get serialized. if size is None: - # split takes a patch count, so round up rather than hand it a float. - size = math.ceil(len(spool) / (os.cpu_count() or 1)) + # split takes a positive patch count, so round up rather than hand + # it a float, and keep an empty spool from asking for zero. + size = max(1, math.ceil(len(spool) / (os.cpu_count() or 1))) spools = list(spool.split(size=size)) # this is a hack to get the progress bar to work. Essentially, we just # add a secret flag to all but one spool so that progress bar is only diff --git a/tests/test_core/test_spool.py b/tests/test_core/test_spool.py index f3d6cf60b..1563b81ab 100644 --- a/tests/test_core/test_spool.py +++ b/tests/test_core/test_spool.py @@ -610,6 +610,13 @@ def test_uneven_size(self, random_spool): assert len(split[0]) == 2 assert len(split[1]) == 1 + @pytest.mark.parametrize("kwargs", [{"size": 0}, {"size": -1}, {"count": 0}]) + def test_non_positive_raises(self, random_spool, kwargs): + """A size or count of zero or less would never finish yielding.""" + msg = "requires a positive size or count" + with pytest.raises(ParameterError, match=msg): + list(random_spool.split(**kwargs)) + def test_non_integral_size(self, random_spool_len_10): """A size which isn't a whole number rounds up rather than raising.""" split = list(random_spool_len_10.split(size=2.5)) diff --git a/tests/test_utils/test_misc.py b/tests/test_utils/test_misc.py index 68b2b8409..48a57cb77 100644 --- a/tests/test_utils/test_misc.py +++ b/tests/test_utils/test_misc.py @@ -15,6 +15,7 @@ import pytest from upath import UPath +import dascore as dc from dascore.exceptions import MissingOptionalDependencyError from dascore.utils.misc import ( _iter_filesystem, @@ -669,6 +670,21 @@ def map(self, func, spools): assert out == [3, 3, 3] assert seen == ["Applying to spool"] + def test_empty_spool_with_client(self): + """An empty spool asks for no work rather than a split size of zero.""" + + class DummyClient: + def map(self, func, spools): + return [func(spool) for spool in spools] + + out = _spool_map( + dc.spool([]), + lambda patch: patch, + client=DummyClient(), + progress=False, + ) + assert out == [] + class Test2DLineIntersection: """Tests for 2D line intersection helper."""