How to use the yapf.yapflib.pytree_utils.SetNodeAnnotation function in yapf

def _SetMustSplitOnFirstLeaf(node):
  """Set the "must split" annotation on the first leaf node."""
  pytree_utils.SetNodeAnnotation(
      pytree_utils.FirstLeafNode(node), pytree_utils.Annotation.MUST_SPLIT,
      True)

Takes care to only increase the penalty. If the node already has a higher
  or equal penalty associated with it, this is a no-op.

  Args:
    tree: subtree to annotate
    annotate_name: name of the annotation to set
    annotate_value: value of the annotation to set
  """
  for child in tree.children:
    _RecAnnotate(child, annotate_name, annotate_value)
  if isinstance(tree, pytree.Leaf):
    cur_annotate = pytree_utils.GetNodeAnnotation(
        tree, annotate_name, default=0)
    if cur_annotate < annotate_value:
      pytree_utils.SetNodeAnnotation(tree, annotate_name, annotate_value)

style.Get('BLANK_LINE_BEFORE_MODULE_DOCSTRING')):
      # Enforce a blank line before a module's docstring.
      return ONE_BLANK_LINE
    # The docstring shouldn't have a newline before it.
    return NO_BLANK_LINES

  prev_last_token = prev_uwline.last
  if prev_last_token.is_docstring:
    if (not indent_depth and first_token.value in {'class', 'def', 'async'}):
      # Separate a class or function from the module-level docstring with
      # appropriate number of blank lines.
      return 1 + style.Get('BLANK_LINES_AROUND_TOP_LEVEL_DEFINITION')
    if (first_nested and
        not style.Get('BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF') and
        _IsClassOrDef(first_token)):
      pytree_utils.SetNodeAnnotation(first_token.node,
                                     pytree_utils.Annotation.NEWLINES, None)
      return NO_BLANK_LINES
    if _NoBlankLinesBeforeCurrentToken(prev_last_token.value, first_token,
                                       prev_last_token):
      return NO_BLANK_LINES
    else:
      return ONE_BLANK_LINE

  if _IsClassOrDef(first_token):
    # TODO(morbo): This can go once the blank line calculator is more
    # sophisticated.
    if not indent_depth:
      # This is a top-level class or function.
      is_inline_comment = prev_last_token.whitespace_prefix.count('\n') == 0
      if (not prev_uwline.disable and prev_last_token.is_comment and
          not is_inline_comment):

def _SetNumNewlines(self, node, num_newlines):
    pytree_utils.SetNodeAnnotation(node, pytree_utils.Annotation.NEWLINES,
                                   num_newlines)

def _AnnotateIndents(tree):
  """Annotate the tree with child_indent annotations.

  A child_indent annotation on a node specifies the indentation (as a string,
  like "  ") of its children. It is inferred from the INDENT child of a node.

  Arguments:
    tree: root of a pytree. The pytree is modified to add annotations to nodes.

  Raises:
    RuntimeError: if the tree is malformed.
  """
  # Annotate the root of the tree with zero indent.
  if tree.parent is None:
    pytree_utils.SetNodeAnnotation(tree, pytree_utils.Annotation.CHILD_INDENT,
                                   '')
  for child in tree.children:
    if child.type == token.INDENT:
      child_indent = pytree_utils.GetNodeAnnotation(
          tree, pytree_utils.Annotation.CHILD_INDENT)
      if child_indent is not None and child_indent != child.value:
        raise RuntimeError('inconsistent indentation for child', (tree, child))
      pytree_utils.SetNodeAnnotation(tree, pytree_utils.Annotation.CHILD_INDENT,
                                     child.value)
    _AnnotateIndents(child)

def _AdjustSplitPenalty(uwline):
  """Visit the node and adjust the split penalties if needed.

  A token shouldn't be split if it's not within a bracket pair. Mark any token
  that's not within a bracket pair as "unbreakable".

  Arguments:
    uwline: (UnwrappedLine) An unwrapped line.
  """
  bracket_level = 0
  for index, token in enumerate(uwline.tokens):
    if index and not bracket_level:
      pytree_utils.SetNodeAnnotation(token.node,
                                     pytree_utils.Annotation.SPLIT_PENALTY,
                                     split_penalty.UNBREAKABLE)
    if token.value in pytree_utils.OPENING_BRACKETS:
      bracket_level += 1
    elif token.value in pytree_utils.CLOSING_BRACKETS:
      bracket_level -= 1