Aluode/PerceptionLabPortable

1import warnings
2from itertools import chain
3
4import pytest
5
6from sklearn import config_context
7from sklearn.utils._chunking import gen_even_slices, get_chunk_n_rows
8from sklearn.utils._testing import assert_array_equal
9
10
11def test_gen_even_slices():
12    # check that gen_even_slices contains all samples
13    some_range = range(10)
14    joined_range = list(chain(*[some_range[slice] for slice in gen_even_slices(10, 3)]))
15    assert_array_equal(some_range, joined_range)
16
17
18@pytest.mark.parametrize(
19    ("row_bytes", "max_n_rows", "working_memory", "expected"),
20    [
21        (1024, None, 1, 1024),
22        (1024, None, 0.99999999, 1023),
23        (1023, None, 1, 1025),
24        (1025, None, 1, 1023),
25        (1024, None, 2, 2048),
26        (1024, 7, 1, 7),
27        (1024 * 1024, None, 1, 1),
28    ],
29)
30def test_get_chunk_n_rows(row_bytes, max_n_rows, working_memory, expected):
31    with warnings.catch_warnings():
32        warnings.simplefilter("error", UserWarning)
33        actual = get_chunk_n_rows(
34            row_bytes=row_bytes,
35            max_n_rows=max_n_rows,
36            working_memory=working_memory,
37        )
38
39    assert actual == expected
40    assert type(actual) is type(expected)
41    with config_context(working_memory=working_memory):
42        with warnings.catch_warnings():
43            warnings.simplefilter("error", UserWarning)
44            actual = get_chunk_n_rows(row_bytes=row_bytes, max_n_rows=max_n_rows)
45        assert actual == expected
46        assert type(actual) is type(expected)
47
48
49def test_get_chunk_n_rows_warns():
50    """Check that warning is raised when working_memory is too low."""
51    row_bytes = 1024 * 1024 + 1
52    max_n_rows = None
53    working_memory = 1
54    expected = 1
55
56    warn_msg = (
57        "Could not adhere to working_memory config. Currently 1MiB, 2MiB required."
58    )
59    with pytest.warns(UserWarning, match=warn_msg):
60        actual = get_chunk_n_rows(
61            row_bytes=row_bytes,
62            max_n_rows=max_n_rows,
63            working_memory=working_memory,
64        )
65
66    assert actual == expected
67    assert type(actual) is type(expected)
68
69    with config_context(working_memory=working_memory):
70        with pytest.warns(UserWarning, match=warn_msg):
71            actual = get_chunk_n_rows(row_bytes=row_bytes, max_n_rows=max_n_rows)
72        assert actual == expected
73        assert type(actual) is type(expected)
74