from pandas import (
    Categorical,
    Index,
    Series,
)
from pandas.core.arrays import BaseMaskedArray


def get_array(obj, col=None):
    """
    Helper method to get array for a DataFrame column or a Series.

    Equivalent of df[col].values, but without going through normal getitem,
    which triggers tracking references / CoW (and we might be testing that
    this is done by some other operation).
    """
    if isinstance(obj, Index):
        arr = obj._values
    elif isinstance(obj, Series) and (col is None or obj.name == col):
        arr = obj._values
    else:
        assert col is not None
        icol = obj.columns.get_loc(col)
        assert isinstance(icol, int)
        arr = obj._get_column_array(icol)
    if isinstance(arr, BaseMaskedArray):
        return arr._data
    elif isinstance(arr, Categorical):
        return arr
    return getattr(arr, "_ndarray", arr)