Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/indexes/extension.py: 59%

1""" 

2Shared methods for Index subclasses backed by ExtensionArray. 

3""" 

4 

5from __future__ import annotations 

6 

7from inspect import signature 

8from typing import ( 

9 TYPE_CHECKING, 

10 TypeVar, 

11) 

12 

13from pandas.util._decorators import cache_readonly 

14 

15from pandas.core.dtypes.generic import ABCDataFrame 

16 

17from pandas.core.indexes.base import Index 

18 

19if TYPE_CHECKING: 

20 from collections.abc import Callable 

21 

22 import numpy as np 

23 

24 from pandas._typing import ( 

25 ArrayLike, 

26 npt, 

27 ) 

28 

29 from pandas.core.arrays import IntervalArray 

30 from pandas.core.arrays._mixins import NDArrayBackedExtensionArray 

31 

32_ExtensionIndexT = TypeVar("_ExtensionIndexT", bound="ExtensionIndex") 

33 

34 

35def _inherit_from_data( 

36 name: str, delegate: type, cache: bool = False, wrap: bool = False 

37): 

38 """ 

39 Make an alias for a method of the underlying ExtensionArray. 

40 

41 Parameters 

42 ---------- 

43 name : str 

44 Name of an attribute the class should inherit from its EA parent. 

45 delegate : class 

46 cache : bool, default False 

47 Whether to convert wrapped properties into cache_readonly 

48 wrap : bool, default False 

49 Whether to wrap the inherited result in an Index. 

50 

51 Returns 

52 ------- 

53 attribute, method, property, or cache_readonly 

54 """ 

55 attr = getattr(delegate, name) 

56 

57 if isinstance(attr, property) or type(attr).__name__ == "getset_descriptor": 

58 # getset_descriptor i.e. property defined in cython class 

59 if cache: 

60 

61 def cached(self): 

62 return getattr(self._data, name) 

63 

64 cached.__name__ = name 

65 cached.__doc__ = attr.__doc__ 

66 method = cache_readonly(cached) 

67 

68 else: 

69 

70 def fget(self): 

71 result = getattr(self._data, name) 

72 if wrap: 

73 if isinstance(result, type(self._data)): 

74 return type(self)._simple_new(result, name=self.name) 

75 elif isinstance(result, ABCDataFrame): 

76 return result.set_index(self) 

77 return Index(result, name=self.name, dtype=result.dtype, copy=False) 

78 return result 

79 

80 def fset(self, value) -> None: 

81 setattr(self._data, name, value) 

82 

83 fget.__name__ = name 

84 fget.__doc__ = attr.__doc__ 

85 

86 method = property(fget, fset) 

87 

88 elif not callable(attr): 

89 # just a normal attribute, no wrapping 

90 method = attr 

91 

92 else: 

93 # error: Incompatible redefinition (redefinition with type "Callable[[Any, 

94 # VarArg(Any), KwArg(Any)], Any]", original type "property") 

95 def method(self, *args, **kwargs): # type: ignore[misc] 

96 if "inplace" in kwargs: 

97 raise ValueError(f"cannot use inplace with {type(self).__name__}") 

98 result = attr(self._data, *args, **kwargs) 

99 if wrap: 

100 if isinstance(result, type(self._data)): 

101 return type(self)._simple_new(result, name=self.name) 

102 elif isinstance(result, ABCDataFrame): 

103 return result.set_index(self) 

104 return Index(result, name=self.name, dtype=result.dtype, copy=False) 

105 return result 

106 

107 # error: "property" has no attribute "__name__" 

108 method.__name__ = name # type: ignore[attr-defined] 

109 method.__doc__ = attr.__doc__ 

110 method.__signature__ = signature(attr) # type: ignore[attr-defined] 

111 return method 

112 

113 

114def inherit_names( 

115 names: list[str], delegate: type, cache: bool = False, wrap: bool = False 

116) -> Callable[[type[_ExtensionIndexT]], type[_ExtensionIndexT]]: 

117 """ 

118 Class decorator to pin attributes from an ExtensionArray to an Index subclass. 

119 

120 Parameters 

121 ---------- 

122 names : List[str] 

123 delegate : class 

124 cache : bool, default False 

125 wrap : bool, default False 

126 Whether to wrap the inherited result in an Index. 

127 """ 

128 

129 def wrapper(cls: type[_ExtensionIndexT]) -> type[_ExtensionIndexT]: 

130 for name in names: 

131 meth = _inherit_from_data(name, delegate, cache=cache, wrap=wrap) 

132 setattr(cls, name, meth) 

133 

134 return cls 

135 

136 return wrapper 

137 

138 

139class ExtensionIndex(Index): 

140 """ 

141 Index subclass for indexes backed by ExtensionArray. 

142 """ 

143 

144 # The base class already passes through to _data: 

145 # size, __len__, dtype 

146 

147 _data: IntervalArray | NDArrayBackedExtensionArray 

148 

149 # --------------------------------------------------------------------- 

150 

151 def _validate_fill_value(self, value): 

152 """ 

153 Convert value to be insertable to underlying array. 

154 """ 

155 return self._data._validate_setitem_value(value) 

156 

157 @cache_readonly 

158 def _isnan(self) -> npt.NDArray[np.bool_]: 

159 # error: Incompatible return value type (got "ExtensionArray", expected 

160 # "ndarray") 

161 return self._data.isna() # type: ignore[return-value] 

162 

163 

164class NDArrayBackedExtensionIndex(ExtensionIndex): 

165 """ 

166 Index subclass for indexes backed by NDArrayBackedExtensionArray. 

167 """ 

168 

169 _data: NDArrayBackedExtensionArray 

170 

171 def _get_engine_target(self) -> np.ndarray: 

172 return self._data._ndarray 

173 

174 def _from_join_target(self, result: np.ndarray) -> ArrayLike: 

175 assert result.dtype == self._data._ndarray.dtype 

176 return self._data._from_backing_data(result)