Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/pandas/core/indexes/extension.py: 62%

1""" 

2Shared methods for Index subclasses backed by ExtensionArray. 

3""" 

4from __future__ import annotations 

5 

6from typing import ( 

7 TYPE_CHECKING, 

8 Callable, 

9 TypeVar, 

10) 

11 

12from pandas.util._decorators import cache_readonly 

13 

14from pandas.core.dtypes.generic import ABCDataFrame 

15 

16from pandas.core.indexes.base import Index 

17 

18if TYPE_CHECKING: 

19 import numpy as np 

20 

21 from pandas._typing import ( 

22 ArrayLike, 

23 npt, 

24 ) 

25 

26 from pandas.core.arrays import IntervalArray 

27 from pandas.core.arrays._mixins import NDArrayBackedExtensionArray 

28 

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

30 

31 

32def _inherit_from_data( 

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

34): 

35 """ 

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

37 

38 Parameters 

39 ---------- 

40 name : str 

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

42 delegate : class 

43 cache : bool, default False 

44 Whether to convert wrapped properties into cache_readonly 

45 wrap : bool, default False 

46 Whether to wrap the inherited result in an Index. 

47 

48 Returns 

49 ------- 

50 attribute, method, property, or cache_readonly 

51 """ 

52 attr = getattr(delegate, name) 

53 

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

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

56 if cache: 

57 

58 def cached(self): 

59 return getattr(self._data, name) 

60 

61 cached.__name__ = name 

62 cached.__doc__ = attr.__doc__ 

63 method = cache_readonly(cached) 

64 

65 else: 

66 

67 def fget(self): 

68 result = getattr(self._data, name) 

69 if wrap: 

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

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

72 elif isinstance(result, ABCDataFrame): 

73 return result.set_index(self) 

74 return Index(result, name=self.name) 

75 return result 

76 

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

78 setattr(self._data, name, value) 

79 

80 fget.__name__ = name 

81 fget.__doc__ = attr.__doc__ 

82 

83 method = property(fget, fset) 

84 

85 elif not callable(attr): 

86 # just a normal attribute, no wrapping 

87 method = attr 

88 

89 else: 

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

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

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

93 if "inplace" in kwargs: 

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

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

96 if wrap: 

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

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

99 elif isinstance(result, ABCDataFrame): 

100 return result.set_index(self) 

101 return Index(result, name=self.name) 

102 return result 

103 

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

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

106 method.__doc__ = attr.__doc__ 

107 return method 

108 

109 

110def inherit_names( 

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

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

113 """ 

114 Class decorator to pin attributes from an ExtensionArray to a Index subclass. 

115 

116 Parameters 

117 ---------- 

118 names : List[str] 

119 delegate : class 

120 cache : bool, default False 

121 wrap : bool, default False 

122 Whether to wrap the inherited result in an Index. 

123 """ 

124 

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

126 for name in names: 

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

128 setattr(cls, name, meth) 

129 

130 return cls 

131 

132 return wrapper 

133 

134 

135class ExtensionIndex(Index): 

136 """ 

137 Index subclass for indexes backed by ExtensionArray. 

138 """ 

139 

140 # The base class already passes through to _data: 

141 # size, __len__, dtype 

142 

143 _data: IntervalArray | NDArrayBackedExtensionArray 

144 

145 # --------------------------------------------------------------------- 

146 

147 def _validate_fill_value(self, value): 

148 """ 

149 Convert value to be insertable to underlying array. 

150 """ 

151 return self._data._validate_setitem_value(value) 

152 

153 @cache_readonly 

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

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

156 # "ndarray") 

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

158 

159 

160class NDArrayBackedExtensionIndex(ExtensionIndex): 

161 """ 

162 Index subclass for indexes backed by NDArrayBackedExtensionArray. 

163 """ 

164 

165 _data: NDArrayBackedExtensionArray 

166 

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

168 return self._data._ndarray 

169 

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

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

172 return self._data._from_backing_data(result)