| """ |
| Helpers for sharing tests between DataFrame/Series |
| """ |
| from __future__ import annotations |
|
|
| from typing import TYPE_CHECKING |
|
|
| from pandas import DataFrame |
|
|
| if TYPE_CHECKING: |
| from pandas._typing import DtypeObj |
|
|
|
|
| def get_dtype(obj) -> DtypeObj: |
| if isinstance(obj, DataFrame): |
| |
| return obj.dtypes.iat[0] |
| else: |
| return obj.dtype |
|
|
|
|
| def get_obj(df: DataFrame, klass): |
| """ |
| For sharing tests using frame_or_series, either return the DataFrame |
| unchanged or return it's first column as a Series. |
| """ |
| if klass is DataFrame: |
| return df |
| return df._ixs(0, axis=1) |
|
|
