Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Helpers for sharing tests between DataFrame/Series
3"""
5from __future__ import annotations
7from typing import TYPE_CHECKING
9from pandas import DataFrame
11if TYPE_CHECKING:
12 from pandas._typing import DtypeObj
14 from pandas import Series
17def get_dtype(obj: DataFrame | Series) -> DtypeObj:
18 if isinstance(obj, DataFrame):
19 # Note: we are assuming only one column
20 return obj.dtypes.iat[0]
21 else:
22 return obj.dtype
25def get_obj(df: DataFrame, klass: type) -> DataFrame | Series:
26 """
27 For sharing tests using frame_or_series, either return the DataFrame
28 unchanged or return its first column as a Series.
29 """
30 if klass is DataFrame:
31 return df
32 return df._ixs(0, axis=1)