Spaces:
Running
Running
| from __future__ import annotations | |
| from typing import Any, List, Optional | |
| from jsonpath_ng import parse | |
| from jsonpath_ng.ext import parse as ext_parse | |
| from jsonpath_ng.exceptions import JsonPathLexerError, JsonPathParserError | |
| from app.core.logger import get_logger | |
| logger = get_logger(__name__) | |
| class KeysExtractor: | |
| """Nested JSON extractor using jsonpath-ng. | |
| Supports two modes: | |
| 1. key_names — recursive lookup by key name at any depth | |
| 2. query — extended JSONPath expression with full path output | |
| """ | |
| def __init__( | |
| self, | |
| data: dict | list, | |
| key_names: Optional[List[str]] = None, | |
| result_limit: Optional[int] = 50, | |
| query: Optional[str] = None, | |
| ): | |
| if query is None and not key_names: | |
| raise ValueError("either `query` or `key_names` must be provided") | |
| if key_names is not None: | |
| if not isinstance(key_names, list): | |
| raise TypeError("key_names must be a list") | |
| for kn in key_names: | |
| if not isinstance(kn, str) or not kn: | |
| raise ValueError(f"each key_name must be a non-empty string, got {kn!r}") | |
| if query is not None: | |
| if not isinstance(query, str) or not query.strip(): | |
| raise ValueError("query must be a non-empty string") | |
| if result_limit is not None and result_limit < 1: | |
| raise ValueError("result_limit must be >= 1 or None") | |
| self._data = data | |
| self._key_names = key_names or [] | |
| self._result_limit = result_limit | |
| self._query = query | |
| def extract(self) -> dict[str, Any]: | |
| if self._query is not None: | |
| return {"matches": self._extract_by_query()} | |
| return {k: self._extract_single(k) for k in self._key_names} | |
| def _extract_by_query(self) -> List[dict[str, Any]]: | |
| try: | |
| expr = ext_parse(self._query) | |
| except (JsonPathLexerError, JsonPathParserError) as exc: | |
| raise ValueError(f"invalid jsonpath query {self._query!r}: {exc}") from exc | |
| matches: List[dict[str, Any]] = [ | |
| {"path": str(m.full_path), "value": m.value} | |
| for m in expr.find(self._data) | |
| ] | |
| if self._result_limit is not None: | |
| matches = matches[: self._result_limit] | |
| return matches | |
| def _extract_single(self, key_name: str) -> List[Any]: | |
| try: | |
| expr = _build_jp_expr(key_name) | |
| except ValueError: | |
| return [] | |
| vals = [m.value for m in expr.find(self._data)] | |
| if self._result_limit is not None: | |
| vals = vals[: self._result_limit] | |
| return vals | |
| def _build_jp_expr(key_name: str) -> Any: | |
| if '"' in key_name and "'" in key_name: | |
| raise ValueError(f"key contains both quote types: {key_name!r}") | |
| if '"' in key_name: | |
| expr_str = "$..['" + key_name + "']" | |
| else: | |
| expr_str = '$..["' + key_name + '"]' | |
| try: | |
| return parse(expr_str) | |
| except (JsonPathLexerError, JsonPathParserError): | |
| raise ValueError(f"unable to build jsonpath expression for key: {key_name!r}") | |
| def get_nested_values(data: dict | list, key_name: str) -> List[Any]: | |
| return KeysExtractor(data, key_names=[key_name], result_limit=None).extract()[key_name] | |