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