Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/pandas/core/indexes/extension.py : 47%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Shared methods for Index subclasses backed by ExtensionArray.
3"""
4from typing import List
6import numpy as np
8from pandas.compat.numpy import function as nv
9from pandas.util._decorators import Appender, cache_readonly
11from pandas.core.dtypes.common import (
12 ensure_platform_int,
13 is_dtype_equal,
14 is_object_dtype,
15)
16from pandas.core.dtypes.generic import ABCSeries
18from pandas.core.arrays import ExtensionArray
19from pandas.core.indexers import deprecate_ndim_indexing
20from pandas.core.indexes.base import Index
21from pandas.core.ops import get_op_result_name
24def inherit_from_data(name: str, delegate, cache: bool = False, wrap: bool = False):
25 """
26 Make an alias for a method of the underlying ExtensionArray.
28 Parameters
29 ----------
30 name : str
31 Name of an attribute the class should inherit from its EA parent.
32 delegate : class
33 cache : bool, default False
34 Whether to convert wrapped properties into cache_readonly
35 wrap : bool, default False
36 Whether to wrap the inherited result in an Index.
38 Returns
39 -------
40 attribute, method, property, or cache_readonly
41 """
43 attr = getattr(delegate, name)
45 if isinstance(attr, property):
46 if cache:
48 def cached(self):
49 return getattr(self._data, name)
51 cached.__name__ = name
52 cached.__doc__ = attr.__doc__
53 method = cache_readonly(cached)
55 else:
57 def fget(self):
58 result = getattr(self._data, name)
59 if wrap:
60 if isinstance(result, type(self._data)):
61 return type(self)._simple_new(result, name=self.name)
62 return Index(result, name=self.name)
63 return result
65 def fset(self, value):
66 setattr(self._data, name, value)
68 fget.__name__ = name
69 fget.__doc__ = attr.__doc__
71 method = property(fget, fset)
73 elif not callable(attr):
74 # just a normal attribute, no wrapping
75 method = attr
77 else:
79 def method(self, *args, **kwargs):
80 result = attr(self._data, *args, **kwargs)
81 if wrap:
82 if isinstance(result, type(self._data)):
83 return type(self)._simple_new(result, name=self.name)
84 return Index(result, name=self.name)
85 return result
87 method.__name__ = name
88 method.__doc__ = attr.__doc__
89 return method
92def inherit_names(names: List[str], delegate, cache: bool = False, wrap: bool = False):
93 """
94 Class decorator to pin attributes from an ExtensionArray to a Index subclass.
96 Parameters
97 ----------
98 names : List[str]
99 delegate : class
100 cache : bool, default False
101 wrap : bool, default False
102 Whether to wrap the inherited result in an Index.
103 """
105 def wrapper(cls):
106 for name in names:
107 meth = inherit_from_data(name, delegate, cache=cache, wrap=wrap)
108 setattr(cls, name, meth)
110 return cls
112 return wrapper
115def _make_wrapped_comparison_op(opname):
116 """
117 Create a comparison method that dispatches to ``._data``.
118 """
120 def wrapper(self, other):
121 if isinstance(other, ABCSeries):
122 # the arrays defer to Series for comparison ops but the indexes
123 # don't, so we have to unwrap here.
124 other = other._values
126 other = _maybe_unwrap_index(other)
128 op = getattr(self._data, opname)
129 return op(other)
131 wrapper.__name__ = opname
132 return wrapper
135def make_wrapped_arith_op(opname):
136 def method(self, other):
137 if (
138 isinstance(other, Index)
139 and is_object_dtype(other.dtype)
140 and type(other) is not Index
141 ):
142 # We return NotImplemented for object-dtype index *subclasses* so they have
143 # a chance to implement ops before we unwrap them.
144 # See https://github.com/pandas-dev/pandas/issues/31109
145 return NotImplemented
146 meth = getattr(self._data, opname)
147 result = meth(_maybe_unwrap_index(other))
148 return _wrap_arithmetic_op(self, other, result)
150 method.__name__ = opname
151 return method
154def _wrap_arithmetic_op(self, other, result):
155 if result is NotImplemented:
156 return NotImplemented
158 if isinstance(result, tuple):
159 # divmod, rdivmod
160 assert len(result) == 2
161 return (
162 _wrap_arithmetic_op(self, other, result[0]),
163 _wrap_arithmetic_op(self, other, result[1]),
164 )
166 if not isinstance(result, Index):
167 # Index.__new__ will choose appropriate subclass for dtype
168 result = Index(result)
170 res_name = get_op_result_name(self, other)
171 result.name = res_name
172 return result
175def _maybe_unwrap_index(obj):
176 """
177 If operating against another Index object, we need to unwrap the underlying
178 data before deferring to the DatetimeArray/TimedeltaArray/PeriodArray
179 implementation, otherwise we will incorrectly return NotImplemented.
181 Parameters
182 ----------
183 obj : object
185 Returns
186 -------
187 unwrapped object
188 """
189 if isinstance(obj, Index):
190 return obj._data
191 return obj
194class ExtensionIndex(Index):
195 """
196 Index subclass for indexes backed by ExtensionArray.
197 """
199 _data: ExtensionArray
201 __eq__ = _make_wrapped_comparison_op("__eq__")
202 __ne__ = _make_wrapped_comparison_op("__ne__")
203 __lt__ = _make_wrapped_comparison_op("__lt__")
204 __gt__ = _make_wrapped_comparison_op("__gt__")
205 __le__ = _make_wrapped_comparison_op("__le__")
206 __ge__ = _make_wrapped_comparison_op("__ge__")
208 def __getitem__(self, key):
209 result = self._data[key]
210 if isinstance(result, type(self._data)):
211 return type(self)(result, name=self.name)
213 # Includes cases where we get a 2D ndarray back for MPL compat
214 deprecate_ndim_indexing(result)
215 return result
217 def __iter__(self):
218 return self._data.__iter__()
220 @property
221 def _ndarray_values(self) -> np.ndarray:
222 return self._data._ndarray_values
224 @Appender(Index.dropna.__doc__)
225 def dropna(self, how="any"):
226 if how not in ("any", "all"):
227 raise ValueError(f"invalid how option: {how}")
229 if self.hasnans:
230 return self._shallow_copy(self._data[~self._isnan])
231 return self._shallow_copy()
233 def repeat(self, repeats, axis=None):
234 nv.validate_repeat(tuple(), dict(axis=axis))
235 result = self._data.repeat(repeats, axis=axis)
236 return self._shallow_copy(result)
238 @Appender(Index.take.__doc__)
239 def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs):
240 nv.validate_take(tuple(), kwargs)
241 indices = ensure_platform_int(indices)
243 taken = self._assert_take_fillable(
244 self._data,
245 indices,
246 allow_fill=allow_fill,
247 fill_value=fill_value,
248 na_value=self._na_value,
249 )
250 return type(self)(taken, name=self.name)
252 def unique(self, level=None):
253 if level is not None:
254 self._validate_index_level(level)
256 result = self._data.unique()
257 return self._shallow_copy(result)
259 def _get_unique_index(self, dropna=False):
260 if self.is_unique and not dropna:
261 return self
263 result = self._data.unique()
264 if dropna and self.hasnans:
265 result = result[~result.isna()]
266 return self._shallow_copy(result)
268 @Appender(Index.map.__doc__)
269 def map(self, mapper, na_action=None):
270 # Try to run function on index first, and then on elements of index
271 # Especially important for group-by functionality
272 try:
273 result = mapper(self)
275 # Try to use this result if we can
276 if isinstance(result, np.ndarray):
277 result = Index(result)
279 if not isinstance(result, Index):
280 raise TypeError("The map function must return an Index object")
281 return result
282 except Exception:
283 return self.astype(object).map(mapper)
285 @Appender(Index.astype.__doc__)
286 def astype(self, dtype, copy=True):
287 if is_dtype_equal(self.dtype, dtype) and copy is False:
288 # Ensure that self.astype(self.dtype) is self
289 return self
291 new_values = self._data.astype(dtype, copy=copy)
293 # pass copy=False because any copying will be done in the
294 # _data.astype call above
295 return Index(new_values, dtype=new_values.dtype, name=self.name, copy=False)