updates re: ndarrybackedextensionarray

Author: chalmerlowe

db_dtypes/core.py

@@ -18,6 +18,7 @@
 import pandas
 import pandas.api.extensions
 from pandas.api.types import is_dtype_equal, is_list_like, is_scalar, pandas_dtype
+from pandas.core.arrays import _mixins
 
 from db_dtypes import pandas_backports
 
@@ -42,9 +43,7 @@ def construct_from_string(cls, name: str):
         return cls()
 
 
-class BaseDatetimeArray(
-    pandas_backports.OpsMixin, pandas_backports.NDArrayBackedExtensionArray
-):
+class BaseDatetimeArray(pandas_backports.OpsMixin, _mixins.NDArrayBackedExtensionArray):
     # scalar used to denote NA value inside our self._ndarray, e.g. -1 for
     # Categorical, iNaT for Period. Outside of object dtype, self.isna() should
     # be exactly locations in self._ndarray with _internal_fill_value. See:

db_dtypes/pandas_backports.py

@@ -19,18 +19,13 @@
 the versions in the later versions of pandas.
 """
 
-from typing import Any
-
-import numpy
 import packaging.version
 import pandas
-from pandas.api.types import is_integer
 import pandas.compat.numpy.function
-import pandas.core.nanops
 
 pandas_release = packaging.version.parse(pandas.__version__).release
 
-# Create aliases for private methods in case they move in a future version.
+# # Create aliases for private methods in case they move in a future version.
 nanall = pandas.core.nanops.nanall
 nanany = pandas.core.nanops.nanany
 nanmax = pandas.core.nanops.nanmax
@@ -77,89 +72,3 @@ def import_default(module_name, force=False, default=None):
 class OpsMixin:
     def _cmp_method(self, other, op):  # pragma: NO COVER
         return NotImplemented
-
-
-# TODO: use public API once NDArrayBackedExtensionArray is added to the
-# pandas.
-# See: https://github.com/pandas-dev/pandas/pull/45544
-# Also: Right now none of this is tested in test_pandas_backports.
-# Temporarily marking this as # pragma: NO COVER just to see how this
-# affects unittest coverage
-@import_default("pandas.core.arrays._mixins")
-class NDArrayBackedExtensionArray(
-    pandas.core.arrays.base.ExtensionArray
-):  # pragma: NO COVER
-    def __init__(self, values, dtype):
-        assert isinstance(values, numpy.ndarray)
-        self._ndarray = values
-        self._dtype = dtype
-
-    @classmethod
-    def _from_backing_data(cls, data):
-        return cls(data, data.dtype)
-
-    def __getitem__(self, index):
-        value = self._ndarray[index]
-        if is_integer(index):
-            return self._box_func(value)
-        return self.__class__(value, self._dtype)
-
-    def __setitem__(self, index, value):
-        self._ndarray[index] = self._validate_setitem_value(value)
-
-    def __len__(self):
-        return len(self._ndarray)
-
-    @property
-    def shape(self):
-        return self._ndarray.shape
-
-    @property
-    def ndim(self) -> int:
-        return self._ndarray.ndim
-
-    @property
-    def size(self) -> int:
-        return self._ndarray.size
-
-    @property
-    def nbytes(self) -> int:
-        return self._ndarray.nbytes
-
-    def copy(self):
-        return self[:]
-
-    def repeat(self, n):
-        return self.__class__(self._ndarray.repeat(n), self._dtype)
-
-    def take(
-        self,
-        indices,
-        *,
-        allow_fill: bool = False,
-        fill_value: Any = None,
-        axis: int = 0,
-    ):
-        from pandas.core.algorithms import take
-
-        if allow_fill:
-            fill_value = self._validate_scalar(fill_value)
-
-        new_data = take(
-            self._ndarray,
-            indices,
-            allow_fill=allow_fill,
-            fill_value=fill_value,
-            axis=axis,
-        )
-        return self._from_backing_data(new_data)
-
-    @classmethod
-    def _concat_same_type(cls, to_concat, axis=0):
-        dtypes = {str(x.dtype) for x in to_concat}
-        if len(dtypes) != 1:
-            raise ValueError("to_concat must have the same dtype (tz)", dtypes)
-
-        new_values = [x._ndarray for x in to_concat]
-        new_values = numpy.concatenate(new_values, axis=axis)
-        return to_concat[0]._from_backing_data(new_values)  # type: ignore[arg-type]