compat.py 566 B

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