Repository: pola-rs/polars
File: py-polars/tests/unit/utils/test_utils.py (Python)

Code before:
def test_timedelta_to_int(td: timedelta, time_unit: TimeUnit, expected: int) -> None:
    assert timedelta_to_int(td, time_unit) == expected


def test_estimated_size() -> None:
    s = pl.Series("n", list(range(100)))
    df = s.to_frame()

    for sz in (s.estimated_size(), s.estimated_size("b"), s.estimated_size("bytes")):
        assert sz == df.estimated_size()

    assert s.estimated_size("kb") == (df.estimated_size("b") / 1024)
    assert s.estimated_size("mb") == (df.estimated_size("kb") / 1024)
    assert s.estimated_size("gb") == (df.estimated_size("mb") / 1024)
    assert s.estimated_size("tb") == (df.estimated_size("gb") / 1024)

    with pytest.raises(ValueError):
        s.estimated_size("milkshake")  # type: ignore[arg-type]


def test_estimated_size_25068() -> None:
    df = pl.select(pl.repeat(0, 100000).cast(pl.Int64).reshape((-1, 1)).arr.to_list())

    # write-read clone of df (via parquet)
    b = io.BytesIO()
    df.write_parquet(b, compression="uncompressed")
    df_clone = pl.read_parquet(b)

    assert df.estimated_size() == df_clone.estimated_size()


@pytest.mark.parametrize(
    ("v1", "v2"),
    [
        ("0.16.8", "0.16.7"),
        ("23.0.0", (3, 1000)),
        ((23, 0, 0), "3.1000"),
        (("0", "0", "2beta"), "0.0.1"),
        (("2", "5", "0", "1"), (2, 5, 0)),
    ],
)
def test_parse_version(v1: Any, v2: Any) -> None:
    assert parse_version(v1) > parse_version(v2)
    assert parse_version(v2) < parse_version(v1)


@pytest.mark.slow
def test_in_notebook() -> None:
    # private function, but easier to test this separately and mock it in the callers
    assert not _in_notebook()


Review comment (suggestion):
Instead of going through parquet, we can just slice the list 

`df.slice(5000).estimated_size()`


Code after:
    assert timedelta_to_int(td, time_unit) == expected


def test_estimated_size() -> None:
    s = pl.Series("n", list(range(100)))
    df = s.to_frame()

    for sz in (s.estimated_size(), s.estimated_size("b"), s.estimated_size("bytes")):
        assert sz == df.estimated_size()

    assert s.estimated_size("kb") == (df.estimated_size("b") / 1024)
    assert s.estimated_size("mb") == (df.estimated_size("kb") / 1024)
    assert s.estimated_size("gb") == (df.estimated_size("mb") / 1024)
    assert s.estimated_size("tb") == (df.estimated_size("gb") / 1024)

    with pytest.raises(ValueError):
        s.estimated_size("milkshake")  # type: ignore[arg-type]


def test_estimated_size_sliced_list_25068() -> None:
    df = pl.select(pl.int_range(10000).cast(pl.List(pl.Int64)))

    assert df.slice(5000).estimated_size() / df.estimated_size() <= 0.5


@pytest.mark.parametrize(
    ("v1", "v2"),
    [
        ("0.16.8", "0.16.7"),
        ("23.0.0", (3, 1000)),
        ((23, 0, 0), "3.1000"),
        (("0", "0", "2beta"), "0.0.1"),
        (("2", "5", "0", "1"), (2, 5, 0)),
    ],
)
def test_parse_version(v1: Any, v2: Any) -> None:
    assert parse_version(v1) > parse_version(v2)
    assert parse_version(v2) < parse_version(v1)


@pytest.mark.slow
def test_in_notebook() -> None:
    # private function, but easier to test this separately and mock it in the callers
    assert not _in_notebook()


@pytest.mark.parametrize(
    ("percentiles", "expected", "inject_median"),
    [
        (None, [0.5], True),
        (0.2, [0.2, 0.5], True),