fastexcel/python/tests/test_durations.py at main · ToucanToco/fastexcel · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
from __future__ import annotations

from datetime import date, datetime, timedelta

import fastexcel
import numpy as np
import pandas as pd
import polars as pl
from pandas.testing import assert_frame_equal as pd_assert_frame_equal
from polars.datatypes import DataType as PolarsDataType
from polars.datatypes import Date as PlDate
from polars.datatypes import Datetime as PlDateTime
from polars.datatypes import Duration as PlDuration
from polars.datatypes import Utf8 as PlUtf8
from polars.testing import assert_frame_equal as pl_assert_frame_equal

from .utils import get_expected_pandas_dtype, path_for_fixture


def test_sheet_with_different_time_types() -> None:
    excel_reader = fastexcel.read_excel(path_for_fixture("dates.ods"))
    sheet = excel_reader.load_sheet_by_idx(0)

    pd_df = sheet.to_pandas()
    pl_df = sheet.to_polars()

    ## dtypes
    assert pd_df["date"].dtype == np.dtype("object")
    assert pd_df["datestr"].dtype == get_expected_pandas_dtype("string")
    assert pd_df["time"].dtype == np.dtype("timedelta64[ms]")
    assert pd_df["datetime"].dtype == np.dtype("datetime64[ms]")
    expected_pl_dtypes: dict[str, PolarsDataType] = {
        "date": PlDate(),
        "datestr": PlUtf8(),
        "time": PlDuration(time_unit="ms"),
        "datetime": PlDateTime(time_unit="ms", time_zone=None),
    }
    assert dict(zip(pl_df.columns, pl_df.dtypes)) == expected_pl_dtypes

    ## Contents

    expected_pd = pd.DataFrame(
        {
            "date": [date(2023, 6, 1)],
            "datestr": ["2023-06-01T02:03:04+02:00"],
            "time": pd.Series([pd.to_timedelta("01:02:03")]).astype("timedelta64[ms]"),
            "datetime": pd.Series([pd.to_datetime("2023-06-01 02:03:04")]).astype("datetime64[ms]"),
        }
    )
    expected_pl = pl.DataFrame(
        {
            "date": [date(2023, 6, 1)],
            "datestr": ["2023-06-01T02:03:04+02:00"],
            "time": [timedelta(hours=1, minutes=2, seconds=3)],
            "datetime": [datetime(2023, 6, 1, 2, 3, 4)],
        },
        schema=expected_pl_dtypes,
    )
    pd_assert_frame_equal(pd_df, expected_pd)
    pl_assert_frame_equal(pl_df, expected_pl)


def test_sheet_with_offset_header_row_and_durations() -> None:
    excel_reader = fastexcel.read_excel(path_for_fixture("single-sheet-skip-rows-durations.xlsx"))
    sheet = excel_reader.load_sheet(0, header_row=10)

    pd_df = sheet.to_pandas()
    pl_df = sheet.to_polars()

    assert pd_df["Tot. Time Away From System"].dtype == np.dtype("timedelta64[ms]")
    assert pd_df["Tot. Time Away From System"].tolist() == [
        pd.Timedelta("01:18:43"),
        pd.Timedelta("07:16:51"),
    ]

    assert pl_df["Tot. Time Away From System"].dtype == pl.Duration(time_unit="ms")
    assert pl_df["Tot. Time Away From System"].to_list() == [
        timedelta(hours=1, minutes=18, seconds=43),
        timedelta(hours=7, minutes=16, seconds=51),
    ]