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 TypeVar 

16 

17import numpy as np 

18 

19from pandas._libs.tslibs import Timestamp 

20from pandas.errors import UndefinedVariableError 

21 

22_KT = TypeVar("_KT") 

23_VT = TypeVar("_VT") 

24 

25 

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

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

28 """ 

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

30 

31 Only works when all passed mapping are mutable. 

32 """ 

33 

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

35 for mapping in self.maps: 

36 if key in mapping: 

37 mapping[key] = value 

38 return 

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

40 

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

42 """ 

43 Raises 

44 ------ 

45 KeyError 

46 If `key` doesn't exist. 

47 """ 

48 for mapping in self.maps: 

49 if key in mapping: 

50 del mapping[key] 

51 return 

52 raise KeyError(key) 

53 

54 

55def ensure_scope( 

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

57) -> Scope: 

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

59 return Scope( 

60 level + 1, 

61 global_dict=global_dict, 

62 local_dict=local_dict, 

63 resolvers=resolvers, 

64 target=target, 

65 ) 

66 

67 

68def _replacer(x) -> str: 

69 """ 

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

71 temporary variables with their calling scope's id. 

72 """ 

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

74 # zeros 

75 try: 

76 hexin = ord(x) 

77 except TypeError: 

78 # bytes literals masquerade as ints when iterating in py3 

79 hexin = x 

80 

81 return hex(hexin) 

82 

83 

84def _raw_hex_id(obj) -> str: 

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

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

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

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

89 

90 

91DEFAULT_GLOBALS = { 

92 "Timestamp": Timestamp, 

93 "datetime": datetime.datetime, 

94 "True": True, 

95 "False": False, 

96 "list": list, 

97 "tuple": tuple, 

98 "inf": np.inf, 

99 "Inf": np.inf, 

100} 

101 

102 

103def _get_pretty_string(obj) -> str: 

104 """ 

105 Return a prettier version of obj. 

106 

107 Parameters 

108 ---------- 

109 obj : object 

110 Object to pretty print 

111 

112 Returns 

113 ------- 

114 str 

115 Pretty print object repr 

116 """ 

117 sio = StringIO() 

118 pprint.pprint(obj, stream=sio) 

119 return sio.getvalue() 

120 

121 

122class Scope: 

123 """ 

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

125 and contexts added by pandas. 

126 

127 Parameters 

128 ---------- 

129 level : int 

130 global_dict : dict or None, optional, default None 

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

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

133 target : object 

134 

135 Attributes 

136 ---------- 

137 level : int 

138 scope : DeepChainMap 

139 target : object 

140 temps : dict 

141 """ 

142 

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

144 level: int 

145 scope: DeepChainMap 

146 resolvers: DeepChainMap 

147 temps: dict 

148 

149 def __init__( 

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

151 ) -> None: 

152 self.level = level + 1 

153 

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

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

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

157 self.target = target 

158 

159 if isinstance(local_dict, Scope): 

160 self.scope.update(local_dict.scope) 

161 if local_dict.target is not None: 

162 self.target = local_dict.target 

163 self._update(local_dict.level) 

164 

165 frame = sys._getframe(self.level) 

166 

167 try: 

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

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

170 # numpy array of pandas objects) 

171 scope_global = self.scope.new_child( 

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

173 ) 

174 self.scope = DeepChainMap(scope_global) 

175 if not isinstance(local_dict, Scope): 

176 scope_local = self.scope.new_child( 

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

178 ) 

179 self.scope = DeepChainMap(scope_local) 

180 finally: 

181 del frame 

182 

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

184 if isinstance(local_dict, Scope): 

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

186 self.resolvers = DeepChainMap(*resolvers) 

187 self.temps = {} 

188 

189 def __repr__(self) -> str: 

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

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

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

193 

194 @property 

195 def has_resolvers(self) -> bool: 

196 """ 

197 Return whether we have any extra scope. 

198 

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

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

201 

202 Returns 

203 ------- 

204 hr : bool 

205 """ 

206 return bool(len(self.resolvers)) 

207 

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

209 """ 

210 Resolve a variable name in a possibly local context. 

211 

212 Parameters 

213 ---------- 

214 key : str 

215 A variable name 

216 is_local : bool 

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

218 the '@' symbol) 

219 

220 Returns 

221 ------- 

222 value : object 

223 The value of a particular variable 

224 """ 

225 try: 

226 # only look for locals in outer scope 

227 if is_local: 

228 return self.scope[key] 

229 

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

231 if self.has_resolvers: 

232 return self.resolvers[key] 

233 

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

235 # no resolvers 

236 assert not is_local and not self.has_resolvers 

237 return self.scope[key] 

238 except KeyError: 

239 try: 

240 # last ditch effort we look in temporaries 

241 # these are created when parsing indexing expressions 

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

243 return self.temps[key] 

244 except KeyError as err: 

245 raise UndefinedVariableError(key, is_local) from err 

246 

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

248 """ 

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

250 

251 Parameters 

252 ---------- 

253 old_key : str 

254 Current variable name to replace 

255 new_key : str 

256 New variable name to replace `old_key` with 

257 new_value : object 

258 Value to be replaced along with the possible renaming 

259 """ 

260 if self.has_resolvers: 

261 maps = self.resolvers.maps + self.scope.maps 

262 else: 

263 maps = self.scope.maps 

264 

265 maps.append(self.temps) 

266 

267 for mapping in maps: 

268 if old_key in mapping: 

269 mapping[new_key] = new_value 

270 return 

271 

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

273 """ 

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

275 

276 Parameters 

277 ---------- 

278 stack : list 

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

280 scopes : sequence of strings 

281 A sequence containing valid stack frame attribute names that 

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

283 """ 

284 variables = itertools.product(scopes, stack) 

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

286 try: 

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

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

289 finally: 

290 # won't remove it, but DECREF it 

291 # in Py3 this probably isn't necessary since frame won't be 

292 # scope after the loop 

293 del frame 

294 

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

296 """ 

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

298 

299 Parameters 

300 ---------- 

301 level : int 

302 """ 

303 sl = level + 1 

304 

305 # add sl frames to the scope starting with the 

306 # most distant and overwriting with more current 

307 # makes sure that we can capture variable scope 

308 stack = inspect.stack() 

309 

310 try: 

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

312 finally: 

313 del stack[:], stack 

314 

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

316 """ 

317 Add a temporary variable to the scope. 

318 

319 Parameters 

320 ---------- 

321 value : object 

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

323 

324 Returns 

325 ------- 

326 str 

327 The name of the temporary variable created. 

328 """ 

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

330 

331 # add to inner most scope 

332 assert name not in self.temps 

333 self.temps[name] = value 

334 assert name in self.temps 

335 

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

337 return name 

338 

339 @property 

340 def ntemps(self) -> int: 

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

342 return len(self.temps) 

343 

344 @property 

345 def full_scope(self) -> DeepChainMap: 

346 """ 

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

348 as a mapping. 

349 

350 Returns 

351 ------- 

352 vars : DeepChainMap 

353 All variables in this scope. 

354 """ 

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

356 return DeepChainMap(*maps)