/
test_utils.py
192 lines (156 loc) · 6.38 KB
/
test_utils.py
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
import pytest
import warnings
import json
import numpy as np
import pandas as pd
from .. import infer_vegalite_type, sanitize_dataframe
def test_infer_vegalite_type():
def _check(arr, typ):
assert infer_vegalite_type(arr) == typ
_check(np.arange(5, dtype=float), "quantitative")
_check(np.arange(5, dtype=int), "quantitative")
_check(np.zeros(5, dtype=bool), "nominal")
_check(pd.date_range("2012", "2013"), "temporal")
_check(pd.timedelta_range(365, periods=12), "temporal")
nulled = pd.Series(np.random.randint(10, size=10))
nulled[0] = None
_check(nulled, "quantitative")
_check(["a", "b", "c"], "nominal")
if hasattr(pytest, "warns"): # added in pytest 2.8
with pytest.warns(UserWarning):
_check([], "nominal")
else:
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
_check([], "nominal")
def test_sanitize_dataframe():
# create a dataframe with various types
df = pd.DataFrame(
{
"s": list("abcde"),
"f": np.arange(5, dtype=float),
"i": np.arange(5, dtype=int),
"b": np.array([True, False, True, True, False]),
"d": pd.date_range("2012-01-01", periods=5, freq="H"),
"c": pd.Series(list("ababc"), dtype="category"),
"c2": pd.Series([1, "A", 2.5, "B", None], dtype="category"),
"o": pd.Series([np.array(i) for i in range(5)]),
"p": pd.date_range("2012-01-01", periods=5, freq="H").tz_localize("UTC"),
}
)
# add some nulls
df.iloc[0, df.columns.get_loc("s")] = None
df.iloc[0, df.columns.get_loc("f")] = np.nan
df.iloc[0, df.columns.get_loc("d")] = pd.NaT
df.iloc[0, df.columns.get_loc("o")] = np.array(np.nan)
# JSON serialize. This will fail on non-sanitized dataframes
print(df[["s", "c2"]])
df_clean = sanitize_dataframe(df)
print(df_clean[["s", "c2"]])
print(df_clean[["s", "c2"]].to_dict())
s = json.dumps(df_clean.to_dict(orient="records"))
print(s)
# Re-construct pandas dataframe
df2 = pd.read_json(s)
# Re-order the columns to match df
df2 = df2[df.columns]
# Re-apply original types
for col in df:
if str(df[col].dtype).startswith("datetime"):
# astype(datetime) introduces time-zone issues:
# to_datetime() does not.
utc = isinstance(df[col].dtype, pd.core.dtypes.dtypes.DatetimeTZDtype)
df2[col] = pd.to_datetime(df2[col], utc=utc)
else:
df2[col] = df2[col].astype(df[col].dtype)
# pandas doesn't properly recognize np.array(np.nan), so change it here
df.iloc[0, df.columns.get_loc("o")] = np.nan
assert df.equals(df2)
def test_sanitize_dataframe_colnames():
df = pd.DataFrame(np.arange(12).reshape(4, 3))
# Test that RangeIndex is converted to strings
df = sanitize_dataframe(df)
assert [isinstance(col, str) for col in df.columns]
# Test that non-string columns result in an error
df.columns = [4, "foo", "bar"]
with pytest.raises(ValueError) as err:
sanitize_dataframe(df)
assert str(err.value).startswith("Dataframe contains invalid column name: 4.")
def test_sanitize_dataframe_timedelta():
df = pd.DataFrame({"r": pd.timedelta_range(start="1 day", periods=4)})
with pytest.raises(ValueError) as err:
sanitize_dataframe(df)
assert str(err.value).startswith('Field "r" has type "timedelta')
def test_sanitize_dataframe_infs():
df = pd.DataFrame({"x": [0, 1, 2, np.inf, -np.inf, np.nan]})
df_clean = sanitize_dataframe(df)
assert list(df_clean.dtypes) == [object]
assert list(df_clean["x"]) == [0, 1, 2, None, None, None]
@pytest.mark.skipif(
not hasattr(pd, "Int64Dtype"),
reason="Nullable integers not supported in pandas v{}".format(pd.__version__),
)
def test_sanitize_nullable_integers():
df = pd.DataFrame(
{
"int_np": [1, 2, 3, 4, 5],
"int64": pd.Series([1, 2, 3, None, 5], dtype="UInt8"),
"int64_nan": pd.Series([1, 2, 3, float("nan"), 5], dtype="Int64"),
"float": [1.0, 2.0, 3.0, 4, 5.0],
"float_null": [1, 2, None, 4, 5],
"float_inf": [1, 2, None, 4, (float("inf"))],
}
)
df_clean = sanitize_dataframe(df)
assert {col.dtype.name for _, col in df_clean.items()} == {"object"}
result_python = {col_name: list(col) for col_name, col in df_clean.items()}
assert result_python == {
"int_np": [1, 2, 3, 4, 5],
"int64": [1, 2, 3, None, 5],
"int64_nan": [1, 2, 3, None, 5],
"float": [1.0, 2.0, 3.0, 4.0, 5.0],
"float_null": [1.0, 2.0, None, 4.0, 5.0],
"float_inf": [1.0, 2.0, None, 4.0, None],
}
@pytest.mark.skipif(
not hasattr(pd, "StringDtype"),
reason="dedicated String dtype not supported in pandas v{}".format(pd.__version__),
)
def test_sanitize_string_dtype():
df = pd.DataFrame(
{
"string_object": ["a", "b", "c", "d"],
"string_string": pd.array(["a", "b", "c", "d"], dtype="string"),
"string_object_null": ["a", "b", None, "d"],
"string_string_null": pd.array(["a", "b", None, "d"], dtype="string"),
}
)
df_clean = sanitize_dataframe(df)
assert {col.dtype.name for _, col in df_clean.items()} == {"object"}
result_python = {col_name: list(col) for col_name, col in df_clean.items()}
assert result_python == {
"string_object": ["a", "b", "c", "d"],
"string_string": ["a", "b", "c", "d"],
"string_object_null": ["a", "b", None, "d"],
"string_string_null": ["a", "b", None, "d"],
}
@pytest.mark.skipif(
not hasattr(pd, "BooleanDtype"),
reason="Nullable boolean dtype not supported in pandas v{}".format(pd.__version__),
)
def test_sanitize_boolean_dtype():
df = pd.DataFrame(
{
"bool_none": pd.array([True, False, None], dtype="boolean"),
"none": pd.array([None, None, None], dtype="boolean"),
"bool": pd.array([True, False, True], dtype="boolean"),
}
)
df_clean = sanitize_dataframe(df)
assert {col.dtype.name for _, col in df_clean.items()} == {"object"}
result_python = {col_name: list(col) for col_name, col in df_clean.items()}
assert result_python == {
"bool_none": [True, False, None],
"none": [None, None, None],
"bool": [True, False, True],
}