|
14 | 14 | REQUIRES_SINGLE_FILE = frozenset(["csv", "joblib", "file"])
|
15 | 15 |
|
16 | 16 |
|
17 |
| -def _assert_is_pandas_df(x): |
| 17 | +def _assert_is_pandas_df(x, file_type: str) -> None: |
18 | 18 | import pandas as pd
|
19 | 19 |
|
20 | 20 | if not isinstance(x, pd.DataFrame):
|
21 | 21 | raise NotImplementedError(
|
22 |
| - "Currently only pandas.DataFrame can be saved to a CSV." |
| 22 | + f"Currently only pandas.DataFrame can be saved as type {file_type!r}." |
23 | 23 | )
|
24 | 24 |
|
25 | 25 |
|
@@ -153,26 +153,26 @@ def save_data(obj, fname, type=None, apply_suffix: bool = True) -> "str | Sequen
|
153 | 153 | final_name = f"{fname}{suffix}"
|
154 | 154 |
|
155 | 155 | if type == "csv":
|
156 |
| - _assert_is_pandas_df(obj) |
| 156 | + _assert_is_pandas_df(obj, file_type=type) |
157 | 157 |
|
158 | 158 | obj.to_csv(final_name, index=False)
|
159 | 159 |
|
160 | 160 | elif type == "arrow":
|
161 | 161 | # NOTE: R pins accepts the type arrow, and saves it as feather.
|
162 | 162 | # we allow reading this type, but raise an error for writing.
|
163 |
| - _assert_is_pandas_df(obj) |
| 163 | + _assert_is_pandas_df(obj, file_type=type) |
164 | 164 |
|
165 | 165 | obj.to_feather(final_name)
|
166 | 166 |
|
167 | 167 | elif type == "feather":
|
168 |
| - _assert_is_pandas_df(obj) |
| 168 | + _assert_is_pandas_df(obj, file_type=type) |
169 | 169 |
|
170 | 170 | raise NotImplementedError(
|
171 | 171 | 'Saving data as type "feather" no longer supported. Use type "arrow" instead.'
|
172 | 172 | )
|
173 | 173 |
|
174 | 174 | elif type == "parquet":
|
175 |
| - _assert_is_pandas_df(obj) |
| 175 | + _assert_is_pandas_df(obj, file_type=type) |
176 | 176 |
|
177 | 177 | obj.to_parquet(final_name)
|
178 | 178 |
|
|
0 commit comments