Edits history of script submission #24148 for ' Transform data (helper)'

• python3

  Ap­pro­ved version

  Edit in Windmill

  Embed

  # 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
  

  Submitted by hugo989 5 hours ago