Transform data

# py312

# requirements:

# jmespath==1.1.0

# xmltodict==1.0.4

# wmill==1.851.0

import json

from typing import Literal, Optional

from xml.parsers.expat import ExpatError

import jmespath

import wmill

import xmltodict

from wmill import S3Object

def main(

    data=None,

    expression: str = "@",

    s3_file: Optional[S3Object] = None,

    input_format: Literal["json", "xml"] = "json",

    write_to_s3: bool = False,

):

    """

    Args:

        data: The value to transform, usually previous results, e.g. results.a, or {"tickets": results.a, "users": results.b} to combine several. JSON or XML text is parsed first.

        expression: JMESPath expression applied to the input. @ (the default) returns the parsed input unchanged.

        s3_file: A JSON or XML file in the workspace object storage, read instead of data for large inputs.

        input_format: xml parses XML text into JSON first: attributes become "@name" keys, text next to attributes "#text", repeated elements lists.

        write_to_s3: Write the result as a JSON file to the workspace object storage and return its S3 object instead.

    """

    # Compile first so a bad expression fails before any input is read. Syntax errors

    # repeat the expression with a caret under the failing column.

    query = jmespath.compile(expression or "@")

    source = "data"

    if s3_file and s3_file.get("s3"):

        if data not in (None, "", {}, []):

            raise ValueError("Pass either data or s3_file, not both")

        # The file is still parsed whole in memory (JMESPath needs the full tree):

        # object storage keeps it out of job args and results, not out of RAM.

        data = wmill.load_s3_file(s3_file)

        source = f"s3_file {s3_file['s3']}"

    if input_format == "xml":

        if not isinstance(data, (str, bytes)):

            raise ValueError(f"input_format is xml but {source} is not XML text")

        try:

            data = xmltodict.parse(data)

        except (ExpatError, ValueError) as e:  # ValueError: entity declarations

            raise ValueError(f"{source} is not valid XML: {e}") from None

    elif isinstance(data, (str, bytes)):

        try:

            data = json.loads(data)

        except (json.JSONDecodeError, UnicodeDecodeError) as e:

            raise ValueError(f"{source} is not valid JSON: {e}") from None

    try:

        result = query.search(data)

    except TypeError as e:

        # jmespath raises on a string/number comparison, typically an XML value.

        raise ValueError(

            f"{e}. If a string is compared with a number, convert it with to_number(...)"

        ) from None

    if write_to_s3:

        return wmill.write_s3_file(

            None, json.dumps(result).encode(), content_type="application/json"

        )

    return result