fmeval.data_loaders.jmespath_util

 1import jmespath
 2import logging
 3from typing import Any, List, Dict, Union, Optional
 4from jmespath.exceptions import JMESPathError
 5from jmespath.parser import ParsedResult
 6from fmeval.exceptions import EvalAlgorithmClientError
 7
 8logger = logging.getLogger(__name__)
 9
10
11def compile_jmespath(jmespath_expression: str):
12    """
13    Compiles a JMESPath expression to be used for JSON data.
14    """
15    try:
16        return jmespath.compile(jmespath_expression)
17    except (TypeError, JMESPathError) as e:
18        raise EvalAlgorithmClientError(f"Unable to compile JMESPath {jmespath_expression}") from e
19
20
21def search_jmespath(
22    jmespath_parser: ParsedResult,
23    jmespath_query_type: str,
24    dataset: Union[Dict, List],
25    dataset_name: str,
26) -> Optional[List[Any]]:
27    """Searches a dataset using a JMESPath query.
28
29    :param jmespath_parser: The JMESPath parser, used for parsing model inputs, model outputs,
30        target outputs, or categories.
31    :param jmespath_query_type: Used for error logging. Will always be the `name` attribute
32        of a fmeval.constants.DatasetColumns enumeration.
33    :param dataset: The data to be searched, already deserialized into a dict/list.
34    :param dataset_name: A name associated with the dataset being parsed for logging purposes.
35    :returns: The result of executing the JMESPath query on the dataset.
36    """
37    try:
38        result = jmespath_parser.search(dataset)
39        if result is None:
40            logger.warning(
41                f"Failed to find {jmespath_query_type} columns in dataset `{dataset_name}` using JMESPath "
42                f"query '{jmespath_parser.expression}'."
43            )
44        return result
45    except ValueError:
46        logger.warning(
47            f"Failed to find {jmespath_query_type} columns in dataset `{dataset_name}` using JMESPath query "
48            f"'{jmespath_parser.expression}'."
49        )
50        return None