Skip to content

Commit

Permalink
Tailor _assert_is_pandas_df error - not just CSV but for all file t… (
Browse files Browse the repository at this point in the history
  • Loading branch information
nathanjmcdougall authored Jul 23, 2024
1 parent cda3b33 commit f3ea387
Showing 1 changed file with 6 additions and 6 deletions.
12 changes: 6 additions & 6 deletions pins/drivers.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,12 @@
REQUIRES_SINGLE_FILE = frozenset(["csv", "joblib", "file"])


def _assert_is_pandas_df(x):
def _assert_is_pandas_df(x, file_type: str) -> None:
import pandas as pd

if not isinstance(x, pd.DataFrame):
raise NotImplementedError(
"Currently only pandas.DataFrame can be saved to a CSV."
f"Currently only pandas.DataFrame can be saved as type {file_type!r}."
)


Expand Down Expand Up @@ -153,26 +153,26 @@ def save_data(obj, fname, type=None, apply_suffix: bool = True) -> "str | Sequen
final_name = f"{fname}{suffix}"

if type == "csv":
_assert_is_pandas_df(obj)
_assert_is_pandas_df(obj, file_type=type)

obj.to_csv(final_name, index=False)

elif type == "arrow":
# NOTE: R pins accepts the type arrow, and saves it as feather.
# we allow reading this type, but raise an error for writing.
_assert_is_pandas_df(obj)
_assert_is_pandas_df(obj, file_type=type)

obj.to_feather(final_name)

elif type == "feather":
_assert_is_pandas_df(obj)
_assert_is_pandas_df(obj, file_type=type)

raise NotImplementedError(
'Saving data as type "feather" no longer supported. Use type "arrow" instead.'
)

elif type == "parquet":
_assert_is_pandas_df(obj)
_assert_is_pandas_df(obj, file_type=type)

obj.to_parquet(final_name)

Expand Down

0 comments on commit f3ea387

Please sign in to comment.