sdks/python/apache_beam/yaml/yaml_utils.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import json
import os
import re
import uuid
from collections.abc import Iterable
from collections.abc import Mapping
from typing import Any

import yaml
from yaml import SafeLoader


class SafeLineLoader(SafeLoader):
  """A yaml loader that attaches line information to mappings and strings."""
  class TaggedString(str):
    """A string class to which we can attach metadata.

    This is primarily used to trace a string's origin back to its place in a
    yaml file.
    """
    def __reduce__(self):
      # Pickle as an ordinary string.
      return str, (str(self), )

  def construct_scalar(self, node):
    value = super().construct_scalar(node)
    if isinstance(value, str):
      value = SafeLineLoader.TaggedString(value)
      value._line_ = node.start_mark.line + 1
    return value

  def construct_mapping(self, node, deep=False):
    mapping = super().construct_mapping(node, deep=deep)
    mapping['__line__'] = node.start_mark.line + 1
    mapping['__uuid__'] = self.create_uuid()
    return mapping

  @classmethod
  def create_uuid(cls):
    return str(uuid.uuid4())

  @classmethod
  def strip_metadata(cls, spec, tagged_str=True):
    if isinstance(spec, Mapping):
      return {
          cls.strip_metadata(key, tagged_str): cls.strip_metadata(
              value, tagged_str)
          for (key, value) in spec.items()
          if key not in ('__line__', '__uuid__')
      }
    elif isinstance(spec, Iterable) and not isinstance(spec, (str, bytes)):
      return [cls.strip_metadata(value, tagged_str) for value in spec]
    elif isinstance(spec, SafeLineLoader.TaggedString) and tagged_str:
      return str(spec)
    else:
      return spec

  @staticmethod
  def get_line(obj):
    if isinstance(obj, dict):
      return obj.get('__line__', 'unknown')
    else:
      return getattr(obj, '_line_', 'unknown')


def patch_yaml(original_str: str, updated):
  """Updates a yaml string to match the updated with minimal changes.

  This only changes the portions of original_str that differ between
  original_str and updated in an attempt to preserve comments and formatting.
  """
  if not original_str and updated:
    return yaml.dump(updated, sort_keys=False)

  if original_str[-1] != '\n':
    # Add a trialing newline to avoid having to constantly check this edge case.
    # (It's also a good idea generally...)
    original_str += '\n'

  # The yaml parser returns positions in terms of line and column numbers.
  # Here we construct the mapping between the two.
  line_starts = [0]
  ix = original_str.find('\n')
  while ix != -1:
    line_starts.append(ix + 1)
    ix = original_str.find('\n', ix + 1)

  def pos(line_or_mark, column=0):
    if isinstance(line_or_mark, yaml.Mark):
      line = line_or_mark.line
      column = line_or_mark.column
    else:
      line = line_or_mark
    return line_starts[line] + column

  # Here we define a custom loader with hooks that record where each element is
  # found so we can swap it out appropriately.
  spans = {}

  class SafeMarkLoader(SafeLoader):
    pass

  # We create special subclass types to ensure each returned node is
  # a distinct object.
  marked_types = {}

  def record_yaml_scalar(constructor):
    def wrapper(self, node):
      raw_data = constructor(self, node)
      typ = type(raw_data)
      if typ not in marked_types:
        marked_types[typ] = type(f'Marked_{typ}', (type(raw_data), ), {})
      marked_data = marked_types[typ](raw_data)
      spans[id(marked_data)] = node.start_mark, node.end_mark
      return marked_data

    return wrapper

  SafeMarkLoader.add_constructor(
      'tag:yaml.org,2002:seq',
      record_yaml_scalar(SafeMarkLoader.construct_sequence))
  SafeMarkLoader.add_constructor(
      'tag:yaml.org,2002:map',
      record_yaml_scalar(SafeMarkLoader.construct_mapping))
  for typ in ('bool', 'int', 'float', 'binary', 'timestamp', 'str'):
    SafeMarkLoader.add_constructor(
        f'tag:yaml.org,2002:{typ}',
        record_yaml_scalar(getattr(SafeMarkLoader, f'construct_yaml_{typ}')))

  #  Now load the original yaml using our special parser.
  original = yaml.load(original_str, Loader=SafeMarkLoader)

  # This (recursively) finds the portion of the original string that must
  # be replaced with new content.
  def diff(a: Any, b: Any) -> Iterable[tuple[int, int, str]]:
    if a == b:
      return
    elif (isinstance(a, dict) and isinstance(b, dict) and
          set(a.keys()) == set(b.keys()) and
          all(id(v) in spans for v in a.values())):
      for k, v in a.items():
        yield from diff(v, b[k])
    elif (isinstance(a, list) and isinstance(b, list) and a and b and
          all(id(v) in spans for v in a)):
      # Diff the matching entries.
      for va, vb in zip(a, b):
        yield from diff(va, vb)
      if len(b) < len(a):
        # Remove extra entries
        yield (
            # End of last preserved element.
            pos(spans[id(a[len(b) - 1])][1]),
            # End of last original element.
            pos(spans[id(a[-1])][1]),
            '')
      elif len(b) > len(a):
        # Add extra entries
        list_start, list_end = spans[id(a)]
        start_char = original_str[pos(list_start)]
        if start_char == '[':
          for v in b[len(a):]:
            yield pos(list_end) - 1, pos(list_end) - 1, ', ' + json.dumps(v)
        else:
          assert start_char == '-'
          indent = original_str[pos(list_start.line):pos(list_start)] + '- '
          content = original_str[pos(list_start):pos(list_end)].rstrip()
          actual_end_pos = pos(list_start) + len(content)
          for v in b[len(a):]:
            if isinstance(v, (list, dict)):
              v_str = (
                  yaml.dump(v, sort_keys=False)
                  # Indent.
                  .replace('\n', '\n' + ' ' * len(indent))
                  # Remove blank line indents.
                  .replace(' ' * len(indent) + '\n', '\n').rstrip())
            else:
              v_str = json.dumps(v)
            yield actual_end_pos, actual_end_pos, '\n' + indent + v_str

    else:
      start, end = spans[id(a)]
      indent = original_str[pos(start.line):pos(start)]
      # We strip trailing whitespace as the "end" of an element is often on
      # a subsequent line where the subsequent element actually starts.
      content = original_str[pos(start):pos(end)].rstrip()
      actual_end_pos = pos(start) + len(content)
      trailing = original_str[actual_end_pos:original_str.
                              find('\n', actual_end_pos)]
      if isinstance(b, (list, dict)):
        if indent.strip() in ('', '-') and not trailing.strip():
          # This element wholly occupies its set of lines, so it is safe to use
          # a multi-line yaml representation (appropriately indented).
          yield (
              pos(start),
              actual_end_pos,
              yaml.dump(b, sort_keys=False)
              # Indent.
              .replace('\n', '\n' + ' ' * len(indent))
              # Remove blank line indents.
              .replace(' ' * len(indent) + '\n', '\n').rstrip())
        else:
          # Force flow style.
          yield (
              pos(start),
              actual_end_pos,
              yaml.dump(b, default_flow_style=True, line_break=False).strip())
      elif isinstance(b, str) and re.match('^[A-Za-z0-9_]+$', b):
        # A simple string literal.
        yield pos(start), actual_end_pos, b
      else:
        # A scalar.
        yield pos(start), actual_end_pos, json.dumps(b)

  # Now stick it all together.
  last_end = 0
  content = []
  for start, end, new_content in sorted(diff(original, updated)):
    content.append(original_str[last_end:start])
    content.append(new_content)
    last_end = end
  content.append(original_str[last_end:])
  return ''.join(content)


def locate_data_file(relpath):
  return os.path.join(os.path.dirname(__file__), relpath)