extraction/keys_extractor.py

"""
Nested JSON key-value extractor using jsonpath-ng.

Walks any JSON-like structure (dict / list) and returns every value
attached to a given key, regardless of how deeply nested it is.
"""
from __future__ import annotations

from typing import Any

from jsonpath_ng import parse
from jsonpath_ng.exceptions import JsonPathLexerError, JsonPathParserError

from logger import get_logger

logger = get_logger(__name__)


class Extractor:
    def __init__(
        self,
        data: dict | list,
        key_names: list[str],
        result_limit: int | None = 50,
    ):
        if not isinstance(key_names, list):
            raise TypeError("key_names must be a list")
        if not key_names:
            raise ValueError("key_names must be a non-empty 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 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
        self._result_limit = result_limit

    def extract(self) -> dict[str, list[Any]]:
        return {k: self._extract_single(k) for k in self._key_names}

    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 Extractor(data, [key_name], result_limit=None).extract()[key_name]