diff --git a/src/hydep/manager.py b/src/hydep/manager.py index 3ad0fe9..196be8c 100644 --- a/src/hydep/manager.py +++ b/src/hydep/manager.py @@ -155,7 +155,7 @@ def _validateSubsteps(self, divisions): if divisions <= 0: raise ValueError(f"Divisions must be positive integer, not {divisions}") return FakeSequence(divisions, maxallowed) - if isinstance(divisions, Sequence): + if isinstance(divisions, (numpy.ndarray, Sequence)): if len(divisions) != maxallowed: raise ValueError( "Number of divisions {} not equal to number of time " diff --git a/tests/test_manager.py b/tests/test_manager.py index 4d9372f..a5bbf99 100644 --- a/tests/test_manager.py +++ b/tests/test_manager.py @@ -130,6 +130,15 @@ def test_managerSubsteps(simpleChain): divset = {safeargs.divisions * x for x in range(1, ntimesteps + 1)} assert len(divset) == ntimesteps + # Test that numpy arrays are good + arraySubs = hydep.Manager( + safeargs.chain, + safeargs.timesteps, + safeargs.power, + numpy.arange(1, ntimesteps + 1), + ) + assert arraySubs.substeps == tuple(range(1, ntimesteps + 1)) + with pytest.raises(TypeError): hydep.Manager(safeargs.chain, safeargs.timesteps, safeargs.power, divset)