Alle Dateien aus dem Pythonkurs
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

86 lines
2.3 KiB

"""
Test extension array that has custom attribute information (not stored on the dtype).
"""
from __future__ import annotations
import numbers
from typing import TYPE_CHECKING
import numpy as np
from pandas.core.dtypes.base import ExtensionDtype
import pandas as pd
from pandas.core.arrays import ExtensionArray
if TYPE_CHECKING:
from pandas._typing import type_t
class FloatAttrDtype(ExtensionDtype):
type = float
name = "float_attr"
na_value = np.nan
@classmethod
def construct_array_type(cls) -> type_t[FloatAttrArray]:
"""
Return the array type associated with this dtype.
Returns
-------
type
"""
return FloatAttrArray
class FloatAttrArray(ExtensionArray):
dtype = FloatAttrDtype()
__array_priority__ = 1000
def __init__(self, values, attr=None) -> None:
if not isinstance(values, np.ndarray):
raise TypeError("Need to pass a numpy array of float64 dtype as values")
if not values.dtype == "float64":
raise TypeError("Need to pass a numpy array of float64 dtype as values")
self.data = values
self.attr = attr
@classmethod
def _from_sequence(cls, scalars, dtype=None, copy=False):
data = np.array(scalars, dtype="float64", copy=copy)
return cls(data)
def __getitem__(self, item):
if isinstance(item, numbers.Integral):
return self.data[item]
else:
# slice, list-like, mask
item = pd.api.indexers.check_array_indexer(self, item)
return type(self)(self.data[item], self.attr)
def __len__(self) -> int:
return len(self.data)
def isna(self):
return np.isnan(self.data)
def take(self, indexer, allow_fill=False, fill_value=None):
from pandas.api.extensions import take
data = self.data
if allow_fill and fill_value is None:
fill_value = self.dtype.na_value
result = take(data, indexer, fill_value=fill_value, allow_fill=allow_fill)
return type(self)(result, self.attr)
def copy(self):
return type(self)(self.data.copy(), self.attr)
@classmethod
def _concat_same_type(cls, to_concat):
data = np.concatenate([x.data for x in to_concat])
attr = to_concat[0].attr if len(to_concat) else None
return cls(data, attr)