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

1""" 

2Module for scope operations 

3""" 

4 

5from __future__ import annotations 

6 

7from collections import ChainMap 

8import datetime 

9import inspect 

10from io import StringIO 

11import itertools 

12import pprint 

13import struct 

14import sys 

15from typing import ( 

16 Any, 

17 TypeVar, 

18) 

19 

20import numpy as np 

21 

22from pandas._libs.tslibs import Timestamp 

23from pandas.errors import UndefinedVariableError 

24 

25_KT = TypeVar("_KT") 

26_VT = TypeVar("_VT") 

27 

28 

29# https://docs.python.org/3/library/collections.html#chainmap-examples-and-recipes 

30class DeepChainMap(ChainMap[_KT, _VT]): 

31 """ 

32 Variant of ChainMap that allows direct updates to inner scopes. 

33 

34 Only works when all passed mapping are mutable. 

35 """ 

36 

37 def __setitem__(self, key: _KT, value: _VT) -> None: 

38 for mapping in self.maps: 

39 if key in mapping: 

40 mapping[key] = value 

41 return 

42 self.maps[0][key] = value 

43 

44 def __delitem__(self, key: _KT) -> None: 

45 """ 

46 Raises 

47 ------ 

48 KeyError 

49 If `key` doesn't exist. 

50 """ 

51 for mapping in self.maps: 

52 if key in mapping: 

53 del mapping[key] 

54 return 

55 raise KeyError(key) 

56 

57 

58def ensure_scope( 

59 level: int, global_dict=None, local_dict=None, resolvers=(), target=None 

60) -> Scope: 

61 """Ensure that we are grabbing the correct scope.""" 

62 return Scope( 

63 level + 1, 

64 global_dict=global_dict, 

65 local_dict=local_dict, 

66 resolvers=resolvers, 

67 target=target, 

68 ) 

69 

70 

71def _replacer(x) -> str: 

72 """ 

73 Replace a number with its hexadecimal representation. Used to tag 

74 temporary variables with their calling scope's id. 

75 """ 

76 # get the hex repr of the binary char and remove 0x and pad by pad_size 

77 # zeros 

78 try: 

79 hexin = ord(x) 

80 except TypeError: 

81 # bytes literals masquerade as ints when iterating 

82 hexin = x 

83 

84 return hex(hexin) 

85 

86 

87def _raw_hex_id(obj) -> str: 

88 """Return the padded hexadecimal id of ``obj``.""" 

89 # interpret as a pointer since that's what really what id returns 

90 packed = struct.pack("@P", id(obj)) 

91 return "".join([_replacer(x) for x in packed]) 

92 

93 

94DEFAULT_GLOBALS = { 

95 "Timestamp": Timestamp, 

96 "datetime": datetime.datetime, 

97 "True": True, 

98 "False": False, 

99 "list": list, 

100 "tuple": tuple, 

101 "inf": np.inf, 

102 "Inf": np.inf, 

103} 

104 

105 

106def _get_pretty_string(obj) -> str: 

107 """ 

108 Return a prettier version of obj. 

109 

110 Parameters 

111 ---------- 

112 obj : object 

113 Object to pretty print 

114 

115 Returns 

116 ------- 

117 str 

118 Pretty print object repr 

119 """ 

120 sio = StringIO() 

121 pprint.pprint(obj, stream=sio) 

122 return sio.getvalue() 

123 

124 

125class Scope: 

126 """ 

127 Object to hold scope, with a few bells to deal with some custom syntax 

128 and contexts added by pandas. 

129 

130 Parameters 

131 ---------- 

132 level : int 

133 global_dict : dict or None, optional, default None 

134 local_dict : dict or Scope or None, optional, default None 

135 resolvers : list-like or None, optional, default None 

136 target : object 

137 

138 Attributes 

139 ---------- 

140 level : int 

141 scope : DeepChainMap 

142 target : object 

143 temps : dict 

144 """ 

145 

146 __slots__ = ["level", "resolvers", "scope", "target", "temps"] 

147 level: int 

148 scope: DeepChainMap[str, Any] 

149 resolvers: DeepChainMap[str, Any] 

150 temps: dict[str, Any] 

151 

152 def __init__( 

153 self, level: int, global_dict=None, local_dict=None, resolvers=(), target=None 

154 ) -> None: 

155 self.level = level + 1 

156 

157 # shallow copy because we don't want to keep filling this up with what 

158 # was there before if there are multiple calls to Scope/_ensure_scope 

159 self.scope = DeepChainMap(DEFAULT_GLOBALS.copy()) 

160 self.target = target 

161 

162 if isinstance(local_dict, Scope): 

163 self.scope.update(local_dict.scope) 

164 if local_dict.target is not None: 

165 self.target = local_dict.target 

166 self._update(local_dict.level) 

167 

168 frame = sys._getframe(self.level) 

169 

170 try: 

171 # shallow copy here because we don't want to replace what's in 

172 # scope when we align terms (alignment accesses the underlying 

173 # numpy array of pandas objects) 

174 scope_global = self.scope.new_child( 

175 (global_dict if global_dict is not None else frame.f_globals).copy() 

176 ) 

177 self.scope = DeepChainMap(scope_global) 

178 if not isinstance(local_dict, Scope): 

179 scope_local = self.scope.new_child( 

180 (local_dict if local_dict is not None else frame.f_locals).copy() 

181 ) 

182 self.scope = DeepChainMap(scope_local) 

183 finally: 

184 del frame 

185 

186 # assumes that resolvers are going from outermost scope to inner 

187 if isinstance(local_dict, Scope): 

188 resolvers += tuple(local_dict.resolvers.maps) 

189 self.resolvers = DeepChainMap(*resolvers) 

190 self.temps = {} 

191 

192 def __repr__(self) -> str: 

193 scope_keys = _get_pretty_string(list(self.scope.keys())) 

194 res_keys = _get_pretty_string(list(self.resolvers.keys())) 

195 return f"{type(self).__name__}(scope={scope_keys}, resolvers={res_keys})" 

196 

197 @property 

198 def has_resolvers(self) -> bool: 

199 """ 

200 Return whether we have any extra scope. 

201 

202 For example, DataFrames pass Their columns as resolvers during calls to 

203 ``DataFrame.eval()`` and ``DataFrame.query()``. 

204 

205 Returns 

206 ------- 

207 hr : bool 

208 """ 

209 return bool(len(self.resolvers)) 

210 

211 def resolve(self, key: str, is_local: bool): 

212 """ 

213 Resolve a variable name in a possibly local context. 

214 

215 Parameters 

216 ---------- 

217 key : str 

218 A variable name 

219 is_local : bool 

220 Flag indicating whether the variable is local or not (prefixed with 

221 the '@' symbol) 

222 

223 Returns 

224 ------- 

225 value : object 

226 The value of a particular variable 

227 """ 

228 try: 

229 # only look for locals in outer scope 

230 if is_local: 

231 return self.scope[key] 

232 

233 # not a local variable so check in resolvers if we have them 

234 if self.has_resolvers: 

235 return self.resolvers[key] 

236 

237 # if we're here that means that we have no locals and we also have 

238 # no resolvers 

239 assert not is_local and not self.has_resolvers 

240 return self.scope[key] 

241 except KeyError: 

242 try: 

243 # last ditch effort we look in temporaries 

244 # these are created when parsing indexing expressions 

245 # e.g., df[df > 0] 

246 return self.temps[key] 

247 except KeyError as err: 

248 raise UndefinedVariableError(key, is_local) from err 

249 

250 def swapkey(self, old_key: str, new_key: str, new_value=None) -> None: 

251 """ 

252 Replace a variable name, with a potentially new value. 

253 

254 Parameters 

255 ---------- 

256 old_key : str 

257 Current variable name to replace 

258 new_key : str 

259 New variable name to replace `old_key` with 

260 new_value : object 

261 Value to be replaced along with the possible renaming 

262 """ 

263 if self.has_resolvers: 

264 maps = self.resolvers.maps + self.scope.maps 

265 else: 

266 maps = self.scope.maps 

267 

268 maps.append(self.temps) 

269 

270 for mapping in maps: 

271 if old_key in mapping: 

272 mapping[new_key] = new_value 

273 return 

274 

275 def _get_vars(self, stack, scopes: list[str]) -> None: 

276 """ 

277 Get specifically scoped variables from a list of stack frames. 

278 

279 Parameters 

280 ---------- 

281 stack : list 

282 A list of stack frames as returned by ``inspect.stack()`` 

283 scopes : sequence of strings 

284 A sequence containing valid stack frame attribute names that 

285 evaluate to a dictionary. For example, ('locals', 'globals') 

286 """ 

287 variables = itertools.product(scopes, stack) 

288 for scope, (frame, _, _, _, _, _) in variables: 

289 d = getattr(frame, f"f_{scope}") 

290 self.scope = DeepChainMap(self.scope.new_child(d)) 

291 

292 def _update(self, level: int) -> None: 

293 """ 

294 Update the current scope by going back `level` levels. 

295 

296 Parameters 

297 ---------- 

298 level : int 

299 """ 

300 sl = level + 1 

301 

302 # add sl frames to the scope starting with the 

303 # most distant and overwriting with more current 

304 # makes sure that we can capture variable scope 

305 stack = inspect.stack() 

306 

307 try: 

308 self._get_vars(stack[:sl], scopes=["locals"]) 

309 finally: 

310 del stack[:], stack 

311 

312 def add_tmp(self, value) -> str: 

313 """ 

314 Add a temporary variable to the scope. 

315 

316 Parameters 

317 ---------- 

318 value : object 

319 An arbitrary object to be assigned to a temporary variable. 

320 

321 Returns 

322 ------- 

323 str 

324 The name of the temporary variable created. 

325 """ 

326 name = f"{type(value).__name__}_{self.ntemps}_{_raw_hex_id(self)}" 

327 

328 # add to inner most scope 

329 assert name not in self.temps 

330 self.temps[name] = value 

331 assert name in self.temps 

332 

333 # only increment if the variable gets put in the scope 

334 return name 

335 

336 @property 

337 def ntemps(self) -> int: 

338 """The number of temporary variables in this scope""" 

339 return len(self.temps) 

340 

341 @property 

342 def full_scope(self) -> DeepChainMap[str, Any]: 

343 """ 

344 Return the full scope for use with passing to engines transparently 

345 as a mapping. 

346 

347 Returns 

348 ------- 

349 vars : DeepChainMap 

350 All variables in this scope. 

351 """ 

352 maps = [self.temps, *self.resolvers.maps, *self.scope.maps] 

353 return DeepChainMap(*maps)