text: # Copyright 2018 Google LLC
  #
  # Licensed 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.
  """Utilities for manipulating an Anomalies proto."""
  
  from __future__ import absolute_import
  from __future__ import division
  
  from __future__ import print_function
  
  from typing import Iterable, FrozenSet, List, Text, Tuple
  import pyarrow as pa
  from tensorflow_data_validation import types
  from tensorflow_data_validation.utils import io_util
  from google.protobuf import text_format
  from tensorflow_metadata.proto.v0 import anomalies_pb2
  
  # LINT.IfChange
  MULTIPLE_ERRORS_SHORT_DESCRIPTION = 'Multiple errors'
  
  
  def _make_updated_descriptions(
      reasons: List[anomalies_pb2.AnomalyInfo.Reason]) -> Tuple[Text, Text]:
    """Returns descriptions based on the specified reasons."""
    # If we only have one reason, use its descriptions. Alternatively, if the only
    # reasons for the anomaly are of type SCHEMA_NEW_COLUMN, then just use one of
    # them for the description.
    if len(reasons) == 1 or all(
        reason.type == anomalies_pb2.AnomalyInfo.SCHEMA_NEW_COLUMN
        for reason in reasons):
      return (reasons[0].description, reasons[0].short_description)
    else:
      return (' '.join([reason.description for reason in reasons]),
              MULTIPLE_ERRORS_SHORT_DESCRIPTION)
  # LINT.ThenChange(../anomalies/schema_anomalies.cc)
  
  
  def remove_anomaly_types(
      anomalies: anomalies_pb2.Anomalies,
      types_to_remove: FrozenSet['anomalies_pb2.AnomalyInfo.Type']) -> None:
    """Removes the specified types of anomaly reasons from an Anomalies proto.
  
    If all reasons for a given feature's anomalies are removed, the entire feature
    will be removed from the Anomalies proto.
  
    Args:
      anomalies: The Anomalies proto from which to remove anomaly reasons of
        the specified types.
      types_to_remove: A set of the types of reasons to remove.
    """
    features_to_remove = []
    for feature_name, anomaly_info in anomalies.anomaly_info.items():
      retained_reasons = [
          reason for reason in anomaly_info.reason
          if reason.type not in types_to_remove
      ]
  
      # Clear the diff regions entirely since we do not have a way of readily
      # separating the comparisons that are attributable to the retained reasons
      # from those that are attributable to the removed reasons.
      anomaly_info.ClearField('diff_regions')
      anomaly_info.ClearField('reason')
      if retained_reasons:
        # If there are anomaly reasons that are retained, update the anomaly info
        # for the feature to include only those retained reasons.
        anomaly_info.reason.extend(retained_reasons)
        # Replace the description and short_description based on the reasons that
        # are retained.
        (updated_description,
         updated_short_description) = _make_updated_descriptions(retained_reasons)
        anomaly_info.description = updated_description
        anomaly_info.short_description = updated_short_description
  
      else:
        # If there are no anomaly types that are retained for a given feature,
        # remove that feature from the anomaly_info map entirely.
        features_to_remove.append(feature_name)
  
    for feature_name in features_to_remove:
      del anomalies.anomaly_info[feature_name]
  
  
  def get_anomalies_slicer(
      anomalies: anomalies_pb2.Anomalies) -> types.SliceFunction:
    """Returns a SliceFunction that For each anomaly, yields (anomaly, example).
  
    Args:
      anomalies: An Anomalies proto from which to generate the list of slice keys.
    """
    def slice_fn(example: pa.RecordBatch) -> Iterable[types.SlicedRecordBatch]:
      for feature_name, anomaly_info in anomalies.anomaly_info.items():
        for anomaly_reason in anomaly_info.reason:
          yield (feature_name + '_' +
                 anomalies_pb2.AnomalyInfo.Type.Name(anomaly_reason.type),
                 example)
  
    return slice_fn
  
  
  def write_anomalies_text(anomalies: anomalies_pb2.Anomalies,
                           output_path: Text) -> None:
    """Writes the Anomalies proto to a file in text format.
  
    Args:
      anomalies: An Anomalies protocol buffer.
      output_path: File path to which to write the Anomalies proto.
  
    Raises:
      TypeError: If the input Anomalies proto is not of the expected type.
    """
    if not isinstance(anomalies, anomalies_pb2.Anomalies):
      raise TypeError(
          'anomalies is of type %s; should be an Anomalies proto.' %
          type(anomalies).__name__)
  
    anomalies_text = text_format.MessageToString(anomalies)
    io_util.write_string_to_file(output_path, anomalies_text)
  
  
  def load_anomalies_text(input_path: Text) -> anomalies_pb2.Anomalies:
    """Loads the Anomalies proto stored in text format in the input path.
  
    Args:
      input_path: File path from which to load the Anomalies proto.
  
    Returns:
      An Anomalies protocol buffer.
    """
    anomalies = anomalies_pb2.Anomalies()
    anomalies_text = io_util.read_file_to_string(input_path)
    text_format.Parse(anomalies_text, anomalies)
    return anomalies
  
  
  def load_anomalies_binary(input_path: Text) -> anomalies_pb2.Anomalies:
    """Loads the Anomalies proto stored in binary format in the input path.
  
    Args:
      input_path: File path from which to load the Anomalies proto.
  
    Returns:
      An Anomalies protocol buffer.
    """
    anomalies_proto = anomalies_pb2.Anomalies()
  
    anomalies_proto.ParseFromString(io_util.read_file_to_string(
        input_path, binary_mode=True))
  
    return anomalies_proto
  