Skip to content

Commit

Permalink
Refactor _get_df_family to raise an error instead.
Browse files Browse the repository at this point in the history
  • Loading branch information
nathanjmcdougall committed Jul 23, 2024
1 parent 47f43a2 commit f41103e
Showing 1 changed file with 15 additions and 18 deletions.
33 changes: 15 additions & 18 deletions pins/drivers.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
from .errors import PinsInsecureReadError

from typing import Literal, Sequence
from typing_extensions import assert_never

# TODO: move IFileSystem out of boards, to fix circular import
# from .boards import IFileSystem
Expand All @@ -24,7 +23,8 @@ def _assert_is_pandas_df(x, file_type: str) -> None:
)


def _get_df_family(df) -> Literal["unknown", "pandas", "polars"]:
def _get_df_family(df) -> Literal["pandas", "polars"]:
"""Return the type of DataFrame, or raise NotImplementedError if we can't decide."""
try:
import polars as pl
except ModuleNotFoundError:
Expand All @@ -36,16 +36,15 @@ def _get_df_family(df) -> Literal["unknown", "pandas", "polars"]:

is_pandas_df = isinstance(df, pd.DataFrame)

if not is_polars_df and not is_pandas_df:
return "unknown"
if is_polars_df and is_pandas_df: # Hybrid DataFrame type!
return "unknown"
if is_polars_df and is_pandas_df:
raise NotImplementedError(
"Hybrid DataFrames (simultaneously pandas and polars) are not supported."
)
elif is_polars_df:
return "polars"
elif is_pandas_df:
return "pandas"
else:
assert_never(df)
raise NotImplementedError(f"Unrecognized DataFrame type: {type(df)}")


def load_path(meta, path_to_version):
Expand Down Expand Up @@ -234,15 +233,13 @@ def save_data(obj, fname, type=None, apply_suffix: bool = True) -> "str | Sequen


def default_title(obj, name):
df_family = _get_df_family(obj)

if df_family in ("pandas", "polars"):
# TODO(compat): title says CSV rather than data.frame
# see https://github.com/machow/pins-python/issues/5
shape_str = " x ".join(map(str, obj.shape))
return f"{name}: a pinned {shape_str} DataFrame"
elif df_family == "unknown":
try:
_get_df_family(obj)
except NotImplementedError:
obj_name = type(obj).__qualname__
return f"{name}: a pinned {obj_name} object"
else:
assert_never(df_family)

# TODO(compat): title says CSV rather than data.frame
# see https://github.com/machow/pins-python/issues/5
shape_str = " x ".join(map(str, obj.shape))
return f"{name}: a pinned {shape_str} DataFrame"

0 comments on commit f41103e

Please sign in to comment.