scope.py 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302
  1. """
  2. Module for scope operations
  3. """
  4. import datetime
  5. import inspect
  6. import itertools
  7. import pprint
  8. import struct
  9. import sys
  10. import numpy as np
  11. from pandas.compat import DeepChainMap, StringIO, map
  12. import pandas as pd # noqa
  13. from pandas.core.base import StringMixin
  14. import pandas.core.computation as compu
  15. def _ensure_scope(level, global_dict=None, local_dict=None, resolvers=(),
  16. target=None, **kwargs):
  17. """Ensure that we are grabbing the correct scope."""
  18. return Scope(level + 1, global_dict=global_dict, local_dict=local_dict,
  19. resolvers=resolvers, target=target)
  20. def _replacer(x):
  21. """Replace a number with its hexadecimal representation. Used to tag
  22. temporary variables with their calling scope's id.
  23. """
  24. # get the hex repr of the binary char and remove 0x and pad by pad_size
  25. # zeros
  26. try:
  27. hexin = ord(x)
  28. except TypeError:
  29. # bytes literals masquerade as ints when iterating in py3
  30. hexin = x
  31. return hex(hexin)
  32. def _raw_hex_id(obj):
  33. """Return the padded hexadecimal id of ``obj``."""
  34. # interpret as a pointer since that's what really what id returns
  35. packed = struct.pack('@P', id(obj))
  36. return ''.join(map(_replacer, packed))
  37. _DEFAULT_GLOBALS = {
  38. 'Timestamp': pd._libs.tslib.Timestamp,
  39. 'datetime': datetime.datetime,
  40. 'True': True,
  41. 'False': False,
  42. 'list': list,
  43. 'tuple': tuple,
  44. 'inf': np.inf,
  45. 'Inf': np.inf,
  46. }
  47. def _get_pretty_string(obj):
  48. """Return a prettier version of obj
  49. Parameters
  50. ----------
  51. obj : object
  52. Object to pretty print
  53. Returns
  54. -------
  55. s : str
  56. Pretty print object repr
  57. """
  58. sio = StringIO()
  59. pprint.pprint(obj, stream=sio)
  60. return sio.getvalue()
  61. class Scope(StringMixin):
  62. """Object to hold scope, with a few bells to deal with some custom syntax
  63. and contexts added by pandas.
  64. Parameters
  65. ----------
  66. level : int
  67. global_dict : dict or None, optional, default None
  68. local_dict : dict or Scope or None, optional, default None
  69. resolvers : list-like or None, optional, default None
  70. target : object
  71. Attributes
  72. ----------
  73. level : int
  74. scope : DeepChainMap
  75. target : object
  76. temps : dict
  77. """
  78. __slots__ = 'level', 'scope', 'target', 'temps'
  79. def __init__(self, level, global_dict=None, local_dict=None, resolvers=(),
  80. target=None):
  81. self.level = level + 1
  82. # shallow copy because we don't want to keep filling this up with what
  83. # was there before if there are multiple calls to Scope/_ensure_scope
  84. self.scope = DeepChainMap(_DEFAULT_GLOBALS.copy())
  85. self.target = target
  86. if isinstance(local_dict, Scope):
  87. self.scope.update(local_dict.scope)
  88. if local_dict.target is not None:
  89. self.target = local_dict.target
  90. self.update(local_dict.level)
  91. frame = sys._getframe(self.level)
  92. try:
  93. # shallow copy here because we don't want to replace what's in
  94. # scope when we align terms (alignment accesses the underlying
  95. # numpy array of pandas objects)
  96. self.scope = self.scope.new_child((global_dict or
  97. frame.f_globals).copy())
  98. if not isinstance(local_dict, Scope):
  99. self.scope = self.scope.new_child((local_dict or
  100. frame.f_locals).copy())
  101. finally:
  102. del frame
  103. # assumes that resolvers are going from outermost scope to inner
  104. if isinstance(local_dict, Scope):
  105. resolvers += tuple(local_dict.resolvers.maps)
  106. self.resolvers = DeepChainMap(*resolvers)
  107. self.temps = {}
  108. def __unicode__(self):
  109. scope_keys = _get_pretty_string(list(self.scope.keys()))
  110. res_keys = _get_pretty_string(list(self.resolvers.keys()))
  111. unicode_str = '{name}(scope={scope_keys}, resolvers={res_keys})'
  112. return unicode_str.format(name=type(self).__name__,
  113. scope_keys=scope_keys,
  114. res_keys=res_keys)
  115. @property
  116. def has_resolvers(self):
  117. """Return whether we have any extra scope.
  118. For example, DataFrames pass Their columns as resolvers during calls to
  119. ``DataFrame.eval()`` and ``DataFrame.query()``.
  120. Returns
  121. -------
  122. hr : bool
  123. """
  124. return bool(len(self.resolvers))
  125. def resolve(self, key, is_local):
  126. """Resolve a variable name in a possibly local context
  127. Parameters
  128. ----------
  129. key : text_type
  130. A variable name
  131. is_local : bool
  132. Flag indicating whether the variable is local or not (prefixed with
  133. the '@' symbol)
  134. Returns
  135. -------
  136. value : object
  137. The value of a particular variable
  138. """
  139. try:
  140. # only look for locals in outer scope
  141. if is_local:
  142. return self.scope[key]
  143. # not a local variable so check in resolvers if we have them
  144. if self.has_resolvers:
  145. return self.resolvers[key]
  146. # if we're here that means that we have no locals and we also have
  147. # no resolvers
  148. assert not is_local and not self.has_resolvers
  149. return self.scope[key]
  150. except KeyError:
  151. try:
  152. # last ditch effort we look in temporaries
  153. # these are created when parsing indexing expressions
  154. # e.g., df[df > 0]
  155. return self.temps[key]
  156. except KeyError:
  157. raise compu.ops.UndefinedVariableError(key, is_local)
  158. def swapkey(self, old_key, new_key, new_value=None):
  159. """Replace a variable name, with a potentially new value.
  160. Parameters
  161. ----------
  162. old_key : str
  163. Current variable name to replace
  164. new_key : str
  165. New variable name to replace `old_key` with
  166. new_value : object
  167. Value to be replaced along with the possible renaming
  168. """
  169. if self.has_resolvers:
  170. maps = self.resolvers.maps + self.scope.maps
  171. else:
  172. maps = self.scope.maps
  173. maps.append(self.temps)
  174. for mapping in maps:
  175. if old_key in mapping:
  176. mapping[new_key] = new_value
  177. return
  178. def _get_vars(self, stack, scopes):
  179. """Get specifically scoped variables from a list of stack frames.
  180. Parameters
  181. ----------
  182. stack : list
  183. A list of stack frames as returned by ``inspect.stack()``
  184. scopes : sequence of strings
  185. A sequence containing valid stack frame attribute names that
  186. evaluate to a dictionary. For example, ('locals', 'globals')
  187. """
  188. variables = itertools.product(scopes, stack)
  189. for scope, (frame, _, _, _, _, _) in variables:
  190. try:
  191. d = getattr(frame, 'f_' + scope)
  192. self.scope = self.scope.new_child(d)
  193. finally:
  194. # won't remove it, but DECREF it
  195. # in Py3 this probably isn't necessary since frame won't be
  196. # scope after the loop
  197. del frame
  198. def update(self, level):
  199. """Update the current scope by going back `level` levels.
  200. Parameters
  201. ----------
  202. level : int or None, optional, default None
  203. """
  204. sl = level + 1
  205. # add sl frames to the scope starting with the
  206. # most distant and overwriting with more current
  207. # makes sure that we can capture variable scope
  208. stack = inspect.stack()
  209. try:
  210. self._get_vars(stack[:sl], scopes=['locals'])
  211. finally:
  212. del stack[:], stack
  213. def add_tmp(self, value):
  214. """Add a temporary variable to the scope.
  215. Parameters
  216. ----------
  217. value : object
  218. An arbitrary object to be assigned to a temporary variable.
  219. Returns
  220. -------
  221. name : basestring
  222. The name of the temporary variable created.
  223. """
  224. name = '{name}_{num}_{hex_id}'.format(name=type(value).__name__,
  225. num=self.ntemps,
  226. hex_id=_raw_hex_id(self))
  227. # add to inner most scope
  228. assert name not in self.temps
  229. self.temps[name] = value
  230. assert name in self.temps
  231. # only increment if the variable gets put in the scope
  232. return name
  233. @property
  234. def ntemps(self):
  235. """The number of temporary variables in this scope"""
  236. return len(self.temps)
  237. @property
  238. def full_scope(self):
  239. """Return the full scope for use with passing to engines transparently
  240. as a mapping.
  241. Returns
  242. -------
  243. vars : DeepChainMap
  244. All variables in this scope.
  245. """
  246. maps = [self.temps] + self.resolvers.maps + self.scope.maps
  247. return DeepChainMap(*maps)