Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/computation/align.py: 23%

1""" 

2Core eval alignment algorithms. 

3""" 

4 

5from __future__ import annotations 

6 

7from functools import ( 

8 partial, 

9 wraps, 

10) 

11from typing import TYPE_CHECKING 

12import warnings 

13 

14import numpy as np 

15 

16from pandas._config.config import get_option 

17 

18from pandas.errors import PerformanceWarning 

19from pandas.util._exceptions import find_stack_level 

20 

21from pandas.core.dtypes.generic import ( 

22 ABCDataFrame, 

23 ABCSeries, 

24) 

25 

26from pandas.core.base import PandasObject 

27import pandas.core.common as com 

28from pandas.core.computation.common import result_type_many 

29 

30if TYPE_CHECKING: 

31 from collections.abc import ( 

32 Callable, 

33 Sequence, 

34 ) 

35 

36 from pandas._typing import F 

37 

38 from pandas.core.generic import NDFrame 

39 from pandas.core.indexes.api import Index 

40 

41 

42def _align_core_single_unary_op( 

43 term, 

44) -> tuple[partial | type[NDFrame], dict[str, Index] | None]: 

45 typ: partial | type[NDFrame] 

46 axes: dict[str, Index] | None = None 

47 

48 if isinstance(term.value, np.ndarray): 

49 typ = partial(np.asanyarray, dtype=term.value.dtype) 

50 else: 

51 typ = type(term.value) 

52 if hasattr(term.value, "axes"): 

53 axes = _zip_axes_from_type(typ, term.value.axes) 

54 

55 return typ, axes 

56 

57 

58def _zip_axes_from_type( 

59 typ: type[NDFrame], new_axes: Sequence[Index] 

60) -> dict[str, Index]: 

61 return {name: new_axes[i] for i, name in enumerate(typ._AXIS_ORDERS)} 

62 

63 

64def _any_pandas_objects(terms) -> bool: 

65 """ 

66 Check a sequence of terms for instances of PandasObject. 

67 """ 

68 return any(isinstance(term.value, PandasObject) for term in terms) 

69 

70 

71def _filter_special_cases(f) -> Callable[[F], F]: 

72 @wraps(f) 

73 def wrapper(terms): 

74 # single unary operand 

75 if len(terms) == 1: 

76 return _align_core_single_unary_op(terms[0]) 

77 

78 term_values = (term.value for term in terms) 

79 

80 # we don't have any pandas objects 

81 if not _any_pandas_objects(terms): 

82 return result_type_many(*term_values), None 

83 

84 return f(terms) 

85 

86 return wrapper 

87 

88 

89@_filter_special_cases 

90def _align_core(terms): 

91 term_index = [i for i, term in enumerate(terms) if hasattr(term.value, "axes")] 

92 term_dims = [terms[i].value.ndim for i in term_index] 

93 

94 from pandas import Series 

95 

96 ndims = Series(dict(zip(term_index, term_dims, strict=True))) 

97 

98 # initial axes are the axes of the largest-axis'd term 

99 biggest = terms[ndims.idxmax()].value 

100 typ = biggest._constructor 

101 axes = biggest.axes 

102 naxes = len(axes) 

103 gt_than_one_axis = naxes > 1 

104 

105 for value in (terms[i].value for i in term_index): 

106 is_series = isinstance(value, ABCSeries) 

107 is_series_and_gt_one_axis = is_series and gt_than_one_axis 

108 

109 for axis, items in enumerate(value.axes): 

110 if is_series_and_gt_one_axis: 

111 ax, itm = naxes - 1, value.index 

112 else: 

113 ax, itm = axis, items 

114 

115 if not axes[ax].is_(itm): 

116 axes[ax] = axes[ax].union(itm) 

117 

118 for i, ndim in ndims.items(): 

119 for axis, items in zip(range(ndim), axes, strict=False): 

120 ti = terms[i].value 

121 

122 if hasattr(ti, "reindex"): 

123 transpose = isinstance(ti, ABCSeries) and naxes > 1 

124 reindexer = axes[naxes - 1] if transpose else items 

125 

126 term_axis_size = len(ti.axes[axis]) 

127 reindexer_size = len(reindexer) 

128 

129 ordm = np.log10(max(1, abs(reindexer_size - term_axis_size))) 

130 if ( 

131 get_option("performance_warnings") 

132 and ordm >= 1 

133 and reindexer_size >= 10000 

134 ): 

135 w = ( 

136 f"Alignment difference on axis {axis} is larger " 

137 f"than an order of magnitude on term {terms[i].name!r}, " 

138 f"by more than {ordm:.4g}; performance may suffer." 

139 ) 

140 warnings.warn( 

141 w, category=PerformanceWarning, stacklevel=find_stack_level() 

142 ) 

143 

144 obj = ti.reindex(reindexer, axis=axis) 

145 terms[i].update(obj) 

146 

147 terms[i].update(terms[i].value.values) 

148 

149 return typ, _zip_axes_from_type(typ, axes) 

150 

151 

152def align_terms(terms): 

153 """ 

154 Align a set of terms. 

155 """ 

156 try: 

157 # flatten the parse tree (a nested list, really) 

158 terms = list(com.flatten(terms)) 

159 except TypeError: 

160 # can't iterate so it must just be a constant or single variable 

161 if isinstance(terms.value, (ABCSeries, ABCDataFrame)): 

162 typ = type(terms.value) 

163 name = terms.value.name if isinstance(terms.value, ABCSeries) else None 

164 return typ, _zip_axes_from_type(typ, terms.value.axes), name 

165 return np.result_type(terms.type), None, None 

166 

167 # if all resolved variables are numeric scalars 

168 if all(term.is_scalar for term in terms): 

169 return result_type_many(*(term.value for term in terms)).type, None, None 

170 

171 # if all input series have a common name, propagate it to the returned series 

172 names = {term.value.name for term in terms if isinstance(term.value, ABCSeries)} 

173 name = names.pop() if len(names) == 1 else None 

174 

175 # perform the main alignment 

176 typ, axes = _align_core(terms) 

177 return typ, axes, name 

178 

179 

180def reconstruct_object(typ, obj, axes, dtype, name): 

181 """ 

182 Reconstruct an object given its type, raw value, and possibly empty 

183 (None) axes. 

184 

185 Parameters 

186 ---------- 

187 typ : object 

188 A type 

189 obj : object 

190 The value to use in the type constructor 

191 axes : dict 

192 The axes to use to construct the resulting pandas object 

193 

194 Returns 

195 ------- 

196 ret : typ 

197 An object of type ``typ`` with the value `obj` and possible axes 

198 `axes`. 

199 """ 

200 try: 

201 typ = typ.type 

202 except AttributeError: 

203 pass 

204 

205 res_t = np.result_type(obj.dtype, dtype) 

206 

207 if not isinstance(typ, partial) and issubclass(typ, PandasObject): 

208 if name is None: 

209 return typ(obj, dtype=res_t, **axes) 

210 return typ(obj, dtype=res_t, name=name, **axes) 

211 

212 # special case for pathological things like ~True/~False 

213 if hasattr(res_t, "type") and typ == np.bool_ and res_t != np.bool_: 

214 ret_value = res_t.type(obj) 

215 else: 

216 ret_value = res_t.type(obj) 

217 # The condition is to distinguish 0-dim array (returned in case of 

218 # scalar) and 1 element array 

219 # e.g. np.array(0) and np.array([0]) 

220 if ( 

221 len(obj.shape) == 1 

222 and len(obj) == 1 

223 and not isinstance(ret_value, np.ndarray) 

224 ): 

225 ret_value = np.array([ret_value]).astype(res_t) 

226 

227 return ret_value