fastexcel/python/tests/test_tables.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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
from datetime import datetime

import fastexcel
import pandas as pd
import polars as pl
import pytest
from pandas.testing import assert_frame_equal as pd_assert_frame_equal
from polars.testing import assert_frame_equal as pl_assert_frame_equal

from .utils import path_for_fixture


@pytest.mark.parametrize("path", ("sheet-with-tables.xlsx",))
def test_table_names(path: str) -> None:
    excel_reader = fastexcel.read_excel(path_for_fixture(path))
    table_names = excel_reader.table_names()

    assert table_names == ["users"]


@pytest.mark.parametrize("path", ("sheet-with-tables.xlsx",))
def test_table_names_with_sheet_name(path: str) -> None:
    excel_reader = fastexcel.read_excel(path_for_fixture(path))
    table_names = excel_reader.table_names("sheet1")

    assert table_names == ["users"]

    table_names = excel_reader.table_names("sheet2")

    assert table_names == []


@pytest.mark.parametrize("path", ("sheet-with-tables.xlsx",))
def test_load_table(path: str) -> None:
    excel_reader = fastexcel.read_excel(path_for_fixture(path))
    users_tbl = excel_reader.load_table("users")

    assert users_tbl.name == "users"
    assert users_tbl.sheet_name == "sheet1"
    assert users_tbl.specified_dtypes is None
    assert users_tbl.available_columns() == [
        fastexcel.ColumnInfo(
            name="User Id",
            index=0,
            absolute_index=0,
            dtype="float",
            dtype_from="guessed",
            column_name_from="provided",
        ),
        fastexcel.ColumnInfo(
            name="FirstName",
            index=1,
            absolute_index=1,
            dtype="string",
            dtype_from="guessed",
            column_name_from="provided",
        ),
        fastexcel.ColumnInfo(
            name="LastName",
            index=2,
            absolute_index=2,
            dtype="string",
            dtype_from="guessed",
            column_name_from="provided",
        ),
        fastexcel.ColumnInfo(
            name="Date",
            index=3,
            absolute_index=3,
            dtype="datetime",
            dtype_from="guessed",
            column_name_from="provided",
        ),
    ]
    assert users_tbl.total_height == 3
    assert users_tbl.offset == 0
    assert users_tbl.height == 3
    assert users_tbl.width == 4

    expected_pl = pl.DataFrame(
        {
            "User Id": [1.0, 2.0, 5.0],
            "FirstName": ["Peter", "John", "Hans"],
            "LastName": ["Müller", "Meier", "Fricker"],
            "Date": [datetime(2020, 1, 1), datetime(2024, 5, 4), datetime(2025, 2, 1)],
        }
    ).with_columns(pl.col("Date").dt.cast_time_unit("ms"))
    pl_assert_frame_equal(users_tbl.to_polars(), expected_pl)

    expected_pd = pd.DataFrame(
        {
            "User Id": [1.0, 2.0, 5.0],
            "FirstName": ["Peter", "John", "Hans"],
            "LastName": ["Müller", "Meier", "Fricker"],
            "Date": pd.Series(
                [datetime(2020, 1, 1), datetime(2024, 5, 4), datetime(2025, 2, 1)]
            ).astype("datetime64[ms]"),
        }
    )

    pd_assert_frame_equal(users_tbl.to_pandas(), expected_pd)

    table_eager = excel_reader.load_table("users", eager=True)
    pl_df = pl.from_arrow(table_eager)
    assert isinstance(pl_df, pl.DataFrame)
    pl_assert_frame_equal(pl_df, expected_pl)
    pd_assert_frame_equal(table_eager.to_pandas(), expected_pd)