Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""manage PyTables query interface via Expressions"""
3from __future__ import annotations
5import ast
6from decimal import (
7 Decimal,
8 InvalidOperation,
9)
10from functools import partial
11from typing import (
12 TYPE_CHECKING,
13 Any,
14 ClassVar,
15 Self,
16 cast,
17)
19import numpy as np
21from pandas._libs import lib
22from pandas._libs.tslibs import (
23 Timedelta,
24 Timestamp,
25)
26from pandas.errors import UndefinedVariableError
28from pandas.core.dtypes.common import is_list_like
30import pandas.core.common as com
31from pandas.core.computation import (
32 expr,
33 ops,
34 scope as _scope,
35)
36from pandas.core.computation.common import ensure_decoded
37from pandas.core.computation.expr import BaseExprVisitor
38from pandas.core.computation.ops import is_term
39from pandas.core.construction import extract_array
40from pandas.core.indexes.base import Index
42from pandas.io.formats.printing import (
43 pprint_thing,
44 pprint_thing_encoded,
45)
47if TYPE_CHECKING:
48 from pandas._typing import (
49 TimeUnit,
50 npt,
51 )
54class PyTablesScope(_scope.Scope):
55 __slots__ = ("queryables",)
57 queryables: dict[str, Any]
59 def __init__(
60 self,
61 level: int,
62 global_dict=None,
63 local_dict=None,
64 queryables: dict[str, Any] | None = None,
65 ) -> None:
66 super().__init__(level + 1, global_dict=global_dict, local_dict=local_dict)
67 self.queryables = queryables or {}
70class Term(ops.Term):
71 env: PyTablesScope
73 def __new__(cls, name, env, side=None, encoding=None):
74 if isinstance(name, str):
75 klass = cls
76 else:
77 klass = Constant
78 return object.__new__(klass)
80 def __init__(self, name, env: PyTablesScope, side=None, encoding=None) -> None:
81 super().__init__(name, env, side=side, encoding=encoding)
83 def _resolve_name(self):
84 # must be a queryables
85 if self.side == "left":
86 # Note: The behavior of __new__ ensures that self.name is a str here
87 if self.name not in self.env.queryables:
88 raise NameError(f"name {self.name!r} is not defined")
89 return self.name
91 # resolve the rhs (and allow it to be None)
92 try:
93 return self.env.resolve(self.name, is_local=False)
94 except UndefinedVariableError:
95 return self.name
97 # read-only property overwriting read/write property
98 @property # type: ignore[misc]
99 def value(self):
100 return self._value
103class Constant(Term):
104 def __init__(self, name, env: PyTablesScope, side=None, encoding=None) -> None:
105 assert isinstance(env, PyTablesScope), type(env)
106 super().__init__(name, env, side=side, encoding=encoding)
108 def _resolve_name(self):
109 return self._name
112class BinOp(ops.BinOp):
113 _max_selectors = 31
115 op: str
116 queryables: dict[str, Any]
117 condition: str | None
119 def __init__(self, op: str, lhs, rhs, queryables: dict[str, Any], encoding) -> None:
120 super().__init__(op, lhs, rhs)
121 self.queryables = queryables
122 self.encoding = encoding
123 self.condition = None
125 def _disallow_scalar_only_bool_ops(self) -> None:
126 pass
128 def prune(self, klass):
129 def pr(left, right):
130 """create and return a new specialized BinOp from myself"""
131 if left is None:
132 return right
133 elif right is None:
134 return left
136 k = klass
137 if isinstance(left, ConditionBinOp):
138 if isinstance(right, ConditionBinOp):
139 k = JointConditionBinOp
140 elif isinstance(left, k):
141 return left
142 elif isinstance(right, k):
143 return right
145 elif isinstance(left, FilterBinOp):
146 if isinstance(right, FilterBinOp):
147 k = JointFilterBinOp
148 elif isinstance(left, k):
149 return left
150 elif isinstance(right, k):
151 return right
153 return k(
154 self.op, left, right, queryables=self.queryables, encoding=self.encoding
155 ).evaluate()
157 left, right = self.lhs, self.rhs
159 if is_term(left) and is_term(right):
160 res = pr(left.value, right.value)
161 elif not is_term(left) and is_term(right):
162 res = pr(left.prune(klass), right.value)
163 elif is_term(left) and not is_term(right):
164 res = pr(left.value, right.prune(klass))
165 elif not (is_term(left) or is_term(right)):
166 res = pr(left.prune(klass), right.prune(klass))
168 return res
170 def conform(self, rhs):
171 """inplace conform rhs"""
172 if not is_list_like(rhs):
173 rhs = [rhs]
174 if isinstance(rhs, np.ndarray):
175 rhs = rhs.ravel()
176 return rhs
178 @property
179 def is_valid(self) -> bool:
180 """return True if this is a valid field"""
181 return self.lhs in self.queryables
183 @property
184 def is_in_table(self) -> bool:
185 """
186 return True if this is a valid column name for generation (e.g. an
187 actual column in the table)
188 """
189 return self.queryables.get(self.lhs) is not None
191 @property
192 def kind(self):
193 """the kind of my field"""
194 return getattr(self.queryables.get(self.lhs), "kind", None)
196 @property
197 def meta(self):
198 """the meta of my field"""
199 return getattr(self.queryables.get(self.lhs), "meta", None)
201 @property
202 def metadata(self):
203 """the metadata of my field"""
204 return getattr(self.queryables.get(self.lhs), "metadata", None)
206 def generate(self, v) -> str:
207 """create and return the op string for this TermValue"""
208 val = v.tostring(self.encoding)
209 return f"({self.lhs} {self.op} {val})"
211 def convert_value(self, conv_val) -> TermValue:
212 """
213 convert the expression that is in the term to something that is
214 accepted by pytables
215 """
217 def stringify(value):
218 if self.encoding is not None:
219 return pprint_thing_encoded(value, encoding=self.encoding)
220 return pprint_thing(value)
222 kind = ensure_decoded(self.kind)
223 meta = ensure_decoded(self.meta)
224 if kind == "datetime" or (kind and kind.startswith("datetime64")):
225 if isinstance(conv_val, (int, float)):
226 conv_val = stringify(conv_val)
227 conv_val = ensure_decoded(conv_val)
228 unit: TimeUnit = "ns"
229 if "[" in kind:
230 unit = cast("TimeUnit", kind.split("[")[-1][:-1])
231 conv_val = Timestamp(conv_val).as_unit(unit)
232 if conv_val.tz is not None:
233 conv_val = conv_val.tz_convert("UTC")
234 return TermValue(conv_val, conv_val._value, kind)
235 elif kind.startswith("timedelta"):
236 unit = "ns"
237 if "[" in kind:
238 unit = cast("TimeUnit", kind.split("[")[-1][:-1])
239 if isinstance(conv_val, str):
240 conv_val = Timedelta(conv_val)
241 elif lib.is_integer(conv_val) or lib.is_float(conv_val):
242 conv_val = Timedelta(conv_val, unit="s")
243 else:
244 conv_val = Timedelta(conv_val)
245 conv_val = conv_val.as_unit(unit)._value
246 return TermValue(int(conv_val), conv_val, kind)
248 elif meta == "category":
249 metadata = extract_array(self.metadata, extract_numpy=True)
250 result: npt.NDArray[np.intp] | np.intp | int
251 if conv_val not in metadata:
252 result = -1
253 else:
254 # Find the index of the first match of conv_val in metadata
255 result = np.flatnonzero(metadata == conv_val)[0]
256 return TermValue(result, result, "integer")
257 elif kind == "integer":
258 try:
259 v_dec = Decimal(conv_val)
260 except InvalidOperation:
261 # GH 54186
262 # convert v to float to raise float's ValueError
263 float(conv_val)
264 else:
265 conv_val = int(v_dec.to_integral_exact(rounding="ROUND_HALF_EVEN"))
266 return TermValue(conv_val, conv_val, kind)
267 elif kind == "float":
268 conv_val = float(conv_val)
269 return TermValue(conv_val, conv_val, kind)
270 elif kind == "bool":
271 if isinstance(conv_val, str):
272 conv_val = conv_val.strip().lower() not in [
273 "false",
274 "f",
275 "no",
276 "n",
277 "none",
278 "0",
279 "[]",
280 "{}",
281 "",
282 ]
283 else:
284 conv_val = bool(conv_val)
285 return TermValue(conv_val, conv_val, kind)
286 elif isinstance(conv_val, str):
287 # string quoting
288 return TermValue(conv_val, stringify(conv_val), "string")
289 else:
290 raise TypeError(
291 f"Cannot compare {conv_val} of type {type(conv_val)} to {kind} column"
292 )
294 def convert_values(self) -> None:
295 pass
298class FilterBinOp(BinOp):
299 filter: tuple[Any, Any, Index] | None = None
301 def __repr__(self) -> str:
302 if self.filter is None:
303 return "Filter: Not Initialized"
304 return pprint_thing(f"[Filter : [{self.filter[0]}] -> [{self.filter[1]}]")
306 def invert(self) -> Self:
307 """invert the filter"""
308 if self.filter is not None:
309 self.filter = (
310 self.filter[0],
311 self.generate_filter_op(invert=True),
312 self.filter[2],
313 )
314 return self
316 def format(self):
317 """return the actual filter format"""
318 return [self.filter]
320 # error: Signature of "evaluate" incompatible with supertype "BinOp"
321 def evaluate(self) -> Self | None: # type: ignore[override]
322 if not self.is_valid:
323 raise ValueError(f"query term is not valid [{self}]")
325 rhs = self.conform(self.rhs)
326 values = list(rhs)
328 if self.is_in_table:
329 # if too many values to create the expression, use a filter instead
330 if self.op in ["==", "!="] and len(values) > self._max_selectors:
331 filter_op = self.generate_filter_op()
332 self.filter = (self.lhs, filter_op, Index(values))
334 return self
335 return None
337 # equality conditions
338 if self.op in ["==", "!="]:
339 filter_op = self.generate_filter_op()
340 self.filter = (self.lhs, filter_op, Index(values))
342 else:
343 raise TypeError(
344 f"passing a filterable condition to a non-table indexer [{self}]"
345 )
347 return self
349 def generate_filter_op(self, invert: bool = False):
350 if (self.op == "!=" and not invert) or (self.op == "==" and invert):
351 return lambda axis, vals: ~axis.isin(vals)
352 else:
353 return lambda axis, vals: axis.isin(vals)
356class JointFilterBinOp(FilterBinOp):
357 def format(self):
358 raise NotImplementedError("unable to collapse Joint Filters")
360 # error: Signature of "evaluate" incompatible with supertype "BinOp"
361 def evaluate(self) -> Self: # type: ignore[override]
362 return self
365class ConditionBinOp(BinOp):
366 def __repr__(self) -> str:
367 return pprint_thing(f"[Condition : [{self.condition}]]")
369 def invert(self):
370 """invert the condition"""
371 # if self.condition is not None:
372 # self.condition = "~(%s)" % self.condition
373 # return self
374 raise NotImplementedError(
375 "cannot use an invert condition when passing to numexpr"
376 )
378 def format(self):
379 """return the actual ne format"""
380 return self.condition
382 # error: Signature of "evaluate" incompatible with supertype "BinOp"
383 def evaluate(self) -> Self | None: # type: ignore[override]
384 if not self.is_valid:
385 raise ValueError(f"query term is not valid [{self}]")
387 # convert values if we are in the table
388 if not self.is_in_table:
389 return None
391 rhs = self.conform(self.rhs)
392 values = [self.convert_value(v) for v in rhs]
394 # equality conditions
395 if self.op in ["==", "!="]:
396 # too many values to create the expression?
397 if len(values) <= self._max_selectors:
398 vs = [self.generate(v) for v in values]
399 self.condition = f"({' | '.join(vs)})"
401 # use a filter after reading
402 else:
403 return None
404 else:
405 self.condition = self.generate(values[0])
407 return self
410class JointConditionBinOp(ConditionBinOp):
411 # error: Signature of "evaluate" incompatible with supertype "BinOp"
412 def evaluate(self) -> Self: # type: ignore[override]
413 self.condition = f"({self.lhs.condition} {self.op} {self.rhs.condition})"
414 return self
417class UnaryOp(ops.UnaryOp):
418 def prune(self, klass):
419 if self.op != "~":
420 raise NotImplementedError("UnaryOp only support invert type ops")
422 operand = self.operand
423 operand = operand.prune(klass)
425 if operand is not None and (
426 (issubclass(klass, ConditionBinOp) and operand.condition is not None)
427 or (
428 not issubclass(klass, ConditionBinOp)
429 and issubclass(klass, FilterBinOp)
430 and operand.filter is not None
431 )
432 ):
433 return operand.invert()
434 return None
437class PyTablesExprVisitor(BaseExprVisitor):
438 const_type: ClassVar[type[ops.Term]] = Constant
439 term_type: ClassVar[type[Term]] = Term
441 def __init__(self, env, engine, parser, **kwargs) -> None:
442 super().__init__(env, engine, parser)
443 for bin_op in self.binary_ops:
444 bin_node = self.binary_op_nodes_map[bin_op]
445 setattr(
446 self,
447 f"visit_{bin_node}",
448 lambda node, bin_op=bin_op: partial(BinOp, bin_op, **kwargs),
449 )
451 def visit_UnaryOp(self, node, **kwargs) -> ops.Term | UnaryOp | None:
452 if isinstance(node.op, (ast.Not, ast.Invert)):
453 return UnaryOp("~", self.visit(node.operand))
454 elif isinstance(node.op, ast.USub):
455 return self.const_type(-self.visit(node.operand).value, self.env)
456 elif isinstance(node.op, ast.UAdd):
457 raise NotImplementedError("Unary addition not supported")
458 # TODO: return None might never be reached
459 return None
461 def visit_Index(self, node, **kwargs):
462 return self.visit(node.value).value
464 def visit_Assign(self, node, **kwargs):
465 cmpr = ast.Compare(
466 ops=[ast.Eq()], left=node.targets[0], comparators=[node.value]
467 )
468 return self.visit(cmpr)
470 def visit_Subscript(self, node, **kwargs) -> ops.Term:
471 # only allow simple subscripts
473 value = self.visit(node.value)
474 slobj = self.visit(node.slice)
475 try:
476 value = value.value
477 except AttributeError:
478 pass
480 if isinstance(slobj, Term):
481 # In py39 np.ndarray lookups with Term containing int raise
482 slobj = slobj.value
484 try:
485 return self.const_type(value[slobj], self.env)
486 except TypeError as err:
487 raise ValueError(f"cannot subscript {value!r} with {slobj!r}") from err
489 def visit_Attribute(self, node, **kwargs):
490 attr = node.attr
491 value = node.value
493 ctx = type(node.ctx)
494 if ctx == ast.Load:
495 # resolve the value
496 resolved = self.visit(value)
498 # try to get the value to see if we are another expression
499 try:
500 resolved = resolved.value
501 except AttributeError:
502 pass
504 try:
505 return self.term_type(getattr(resolved, attr), self.env)
506 except AttributeError:
507 # something like datetime.datetime where scope is overridden
508 if isinstance(value, ast.Name) and value.id == attr:
509 return resolved
511 raise ValueError(f"Invalid Attribute context {ctx.__name__}")
513 def translate_In(self, op):
514 return ast.Eq() if isinstance(op, ast.In) else op
516 def _rewrite_membership_op(self, node, left, right):
517 return self.visit(node.op), node.op, left, right
520def _validate_where(w):
521 """
522 Validate that the where statement is of the right type.
524 The type may either be String, Expr, or list-like of Exprs.
526 Parameters
527 ----------
528 w : String term expression, Expr, or list-like of Exprs.
530 Returns
531 -------
532 where : The original where clause if the check was successful.
534 Raises
535 ------
536 TypeError : An invalid data type was passed in for w (e.g. dict).
537 """
538 if not (isinstance(w, (PyTablesExpr, str)) or is_list_like(w)):
539 raise TypeError(
540 "where must be passed as a string, PyTablesExpr, "
541 "or list-like of PyTablesExpr"
542 )
544 return w
547class PyTablesExpr(expr.Expr):
548 """
549 Hold a pytables-like expression, comprised of possibly multiple 'terms'.
551 Parameters
552 ----------
553 where : string term expression, PyTablesExpr, or list-like of PyTablesExprs
554 queryables : a "kinds" map (dict of column name -> kind), or None if column
555 is non-indexable
556 encoding : an encoding that will encode the query terms
558 Returns
559 -------
560 a PyTablesExpr object
562 Examples
563 --------
564 'index>=date'
565 "columns=['A', 'D']"
566 'columns=A'
567 'columns==A'
568 "~(columns=['A','B'])"
569 'index>df.index[3] & string="bar"'
570 '(index>df.index[3] & index<=df.index[6]) | string="bar"'
571 "ts>=Timestamp('2012-02-01')"
572 "major_axis>=20130101"
573 """
575 _visitor: PyTablesExprVisitor | None
576 env: PyTablesScope
577 expr: str
579 def __init__(
580 self,
581 where,
582 queryables: dict[str, Any] | None = None,
583 encoding=None,
584 scope_level: int = 0,
585 ) -> None:
586 where = _validate_where(where)
588 self.encoding = encoding
589 self.condition = None
590 self.filter = None
591 self.terms = None
592 self._visitor = None
594 # capture the environment if needed
595 local_dict: _scope.DeepChainMap[Any, Any] | None = None
597 if isinstance(where, PyTablesExpr):
598 local_dict = where.env.scope
599 _where = where.expr
601 elif is_list_like(where):
602 where = list(where)
603 for idx, w in enumerate(where):
604 if isinstance(w, PyTablesExpr):
605 local_dict = w.env.scope
606 else:
607 where[idx] = _validate_where(w)
608 _where = " & ".join([f"({w})" for w in com.flatten(where)])
609 else:
610 # _validate_where ensures we otherwise have a string
611 _where = where
613 self.expr = _where
614 self.env = PyTablesScope(scope_level + 1, local_dict=local_dict)
616 if queryables is not None and isinstance(self.expr, str):
617 self.env.queryables.update(queryables)
618 self._visitor = PyTablesExprVisitor(
619 self.env,
620 queryables=queryables,
621 parser="pytables",
622 engine="pytables",
623 encoding=encoding,
624 )
625 self.terms = self.parse()
627 def __repr__(self) -> str:
628 if self.terms is not None:
629 return pprint_thing(self.terms)
630 return pprint_thing(self.expr)
632 def evaluate(self):
633 """create and return the numexpr condition and filter"""
634 try:
635 self.condition = self.terms.prune(ConditionBinOp)
636 except AttributeError as err:
637 raise ValueError(
638 f"cannot process expression [{self.expr}], [{self}] "
639 "is not a valid condition"
640 ) from err
641 try:
642 self.filter = self.terms.prune(FilterBinOp)
643 except AttributeError as err:
644 raise ValueError(
645 f"cannot process expression [{self.expr}], [{self}] "
646 "is not a valid filter"
647 ) from err
649 return self.condition, self.filter
652class TermValue:
653 """hold a term value the we use to construct a condition/filter"""
655 def __init__(self, value, converted, kind: str) -> None:
656 assert isinstance(kind, str), kind
657 self.value = value
658 self.converted = converted
659 self.kind = kind
661 def tostring(self, encoding) -> str:
662 """quote the string if not encoded else encode and return"""
663 if self.kind == "string":
664 if encoding is not None:
665 return str(self.converted)
666 return f'"{self.converted}"'
667 elif self.kind == "float":
668 # python 2 str(float) is not always
669 # round-trippable so use repr()
670 return repr(self.converted)
671 return str(self.converted)
674def maybe_expression(s) -> bool:
675 """loose checking if s is a pytables-acceptable expression"""
676 if not isinstance(s, str):
677 return False
678 operations = PyTablesExprVisitor.binary_ops + PyTablesExprVisitor.unary_ops + ("=",)
680 # make sure we have an op at least
681 return any(op in s for op in operations)