1"""
2Helpers for sharing tests between DataFrame/Series
3"""
4from __future__ import annotations
5
6from typing import TYPE_CHECKING
7
8from pandas import DataFrame
9
10if TYPE_CHECKING:
11 from pandas._typing import DtypeObj
12
13
14def get_dtype(obj) -> DtypeObj:
15 if isinstance(obj, DataFrame):
16 # Note: we are assuming only one column
17 return obj.dtypes.iat[0]
18 else:
19 return obj.dtype
20
21
22def get_obj(df: DataFrame, klass):
23 """
24 For sharing tests using frame_or_series, either return the DataFrame
25 unchanged or return it's first column as a Series.
26 """
27 if klass is DataFrame:
28 return df
29 return df._ixs(0, axis=1)