arrays.py 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. """
  2. Methods for cleaning, validating, and unboxing arrays.
  3. """
  4. from pandas.core.dtypes.generic import ABCIndexClass, ABCPandasArray, ABCSeries
  5. def extract_array(obj, extract_numpy=False):
  6. """
  7. Extract the ndarray or ExtensionArray from a Series or Index.
  8. For all other types, `obj` is just returned as is.
  9. Parameters
  10. ----------
  11. obj : object
  12. For Series / Index, the underlying ExtensionArray is unboxed.
  13. For Numpy-backed ExtensionArrays, the ndarray is extracted.
  14. extract_numpy : bool, default False
  15. Whether to extract the ndarray from a PandasArray
  16. Returns
  17. -------
  18. arr : object
  19. Examples
  20. --------
  21. >>> extract_array(pd.Series(['a', 'b', 'c'], dtype='category'))
  22. [a, b, c]
  23. Categories (3, object): [a, b, c]
  24. Other objects like lists, arrays, and DataFrames are just passed through.
  25. >>> extract_array([1, 2, 3])
  26. [1, 2, 3]
  27. For an ndarray-backed Series / Index a PandasArray is returned.
  28. >>> extract_array(pd.Series([1, 2, 3]))
  29. <PandasArray>
  30. [1, 2, 3]
  31. Length: 3, dtype: int64
  32. To extract all the way down to the ndarray, pass ``extract_numpy=True``.
  33. >>> extract_array(pd.Series([1, 2, 3]), extract_numpy=True)
  34. array([1, 2, 3])
  35. """
  36. if isinstance(obj, (ABCIndexClass, ABCSeries)):
  37. obj = obj.array
  38. if extract_numpy and isinstance(obj, ABCPandasArray):
  39. obj = obj.to_numpy()
  40. return obj