Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-02 07:18:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

utils: extend to_chunks functionality, can specify number of chunks or length of chunk

Browse Source
This commit is contained in:
pared
2019-06-28 13:52:59 +02:00
committed by pawel
parent 3d5c42e2fe
commit 7f8da34cf2
6 changed files with 56 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/utils/test_utils.py
+23 -9
View File
@@ -1,16 +1,30 @@
import pytest
from dvc.utils import chunk
from dvc.utils import to_chunks
@pytest.mark.parametrize(
"list_to_chunk, chunk_size, expected_chunks",
[
([1, 2, 3, 4], 1, [[1], [2], [3], [4]]),
([1, 2, 3, 4], 2, [[1, 2], [3, 4]]),
([1, 2, 3, 4], 3, [[1, 2, 3], [4]]),
],
"chunk_size, expected_chunks",
[(1, [[1], [2], [3], [4]]), (2, [[1, 2], [3, 4]]), (3, [[1, 2, 3], [4]])],
)
def test_chunk(list_to_chunk, chunk_size, expected_chunks):
result = list(chunk(list_to_chunk, chunk_size))
def test_to_chunks_chunk_size(chunk_size, expected_chunks):
list_to_chunk = [1, 2, 3, 4]
result = list(to_chunks(list_to_chunk, chunk_size=chunk_size))
assert result == expected_chunks
@pytest.mark.parametrize("num_chunks, chunk_size", [(1, 2), (None, None)])
def test_to_chunks_should_raise(num_chunks, chunk_size):
list_to_chunk = [1, 2, 3]
with pytest.raises(ValueError):
to_chunks(list_to_chunk, num_chunks, chunk_size)
@pytest.mark.parametrize(
"num_chunks, expected_chunks",
[(4, [[1], [2], [3], [4]]), (3, [[1, 2], [3, 4]]), (2, [[1, 2], [3, 4]])],
)
def test_to_chunks_num_chunks(num_chunks, expected_chunks):
list_to_chunk = [1, 2, 3, 4]
result = to_chunks(list_to_chunk, num_chunks=num_chunks)
assert result == expected_chunks