# Path: plur/utils/constants.py
# STAGE_1_DIRNAME = 'stage_1'
# RAW_DATA_DIRNAME = 'raw_data'
# GIT_REPO_DIRNAME = 'git_repo'
# GRAPH_TO_OUTPUT_EXAMPLE_DIRNAME = 'graph_to_output_example'
# TRAIN_SPLIT_NAME = 'train'
# VALIDATION_SPLIT_NAME = 'valid'
# TEST_SPLIT_NAME = 'test'
# METADATA_FILENAME = 'dataset_metadata.json'
# STAGE_2_DIRNAME = 'stage_2'
# VOCAB_FILES_DIRNAME = 'vocab_files'
# NODE_TYPE_VOCAB_FILENAME = 'node_type_vocab.txt'
# NODE_LABEL_VOCAB_FILENAME = 'node_label_vocab.txt'
# EDGE_TYPE_VOCAB_FILENAME = 'edge_type_vocab.txt'
# OUTPUT_TOKEN_VOCAB_FILENAME = 'output_token_vocab.txt'
# TFRECORD_DIRNAME = 'tfrecords'
# POINTER_TOKEN = 'POINTER'
# CLASS_TOKEN_PREFIX = 'CLASS_'
# DONE_TOKEN = 'DONE'
# PAD_TOKEN = '<pad>'
# OOV_TOKEN = '<OOV>'
# RESERVED_TOKENS = [OOV_TOKEN, PAD_TOKEN]
# NUM_RESERVED_TOKENS = len(RESERVED_TOKENS)
# EXAMPLE_GUID = 'graph/global_id'  # For future use.
# KEY_PROVENANCE = 'metadata/provenance'
# KEY_OUTPUT_TOKEN_IDS = 'targets/token_ids'
# KEY_OUTPUT_TOKEN_TEXTS = 'targets/texts'
# KEY_COPY_INPUT_INDICES = 'nodes/copy_indices'
# KEY_COPY_OUTPUT_INDICES = 'targets/copy_indices'
# KEY_POINTER_INPUT_INDICES = 'nodes/pointer_indices'
# KEY_POINTER_OUTPUT_INDICES = 'targets/pointer_indices'
# KEY_NODE_TOKEN_IDS = 'nodes/node_id_sequences'
# KEY_NODE_TEXTS = 'nodes/node_text_sequences'
# KEY_NODE_TOKEN_LENGTHS = 'nodes/node_sequence_lengths'
# KEY_NODE_GUIDS = 'nodes/global_node_id'  # For future use.
# KEY_NODE_TYPE_IDS = 'nodes/node_types'
# KEY_EDGE_TYPE_IDS = 'edges/types'
# KEY_EDGE_SOURCE_INDICES = 'edges/source_indices'
# KEY_EDGE_DESTINATION_INDICES = 'edges/dest_indices'
# KEY_MASKING_CANDIDATE_INDICES = 'masking/candidate_indices'
# PROVENANCE_TENSOR_NAME = 'provenance'
# TARGET_TOKEN_IDS_TENSOR_NAME = 'target_token_ids'
# COPY_INDICES_TENSOR_NAME = 'copy_indices'
# POINTER_INDICES_TENSOR_NAME = 'pointer_indices'
# NODE_ID_SEQUENCES_TENSOR_NAME = 'node_id_sequences'
# NODE_TYPES_TENSOR_NAME = 'node_types'
# EDGE_INDICATORS_TENSOR_NAME = 'edge_indicators'
# NODE_TEXT_SEQUENCES_TENSOR_NAME = 'node_text_sequences'
# TARGET_TEXTS_TENSOR_NAME = 'target_texts'
# MASKING_CANDIDATE_TENSOR_NAME = 'candidate_indices'
# HOPPITY_ADD_NODE_OP_NAME = 'add_node'
# HOPPITY_DEL_NODE_OP_NAME = 'del_node'
# HOPPITY_REPLACE_VAL_OP_NAME = 'replace_val'
# HOPPITY_REPLACE_TYPE_OP_NAME = 'replace_type'
# HOPPITY_REPLACE_NOOP_OP_NAME = 'NoOp'
# HOPPITY_OUTPUT_UNKNOWN_TOKEN = 'UNKNOWN'
# NATURALIZE_CONTEXT_TOKEN_NODE_NAME = 'CONTEXT_TOKEN'
# NATURALIZE_NEXT_TOKEN_EDGE_NAME = 'NEXT_TOKEN'
# NATURALIZE_CONTEXT_ROOT_NODE_NAME = 'CONTEXT_ROOT'
# NATURALIZE_CONTEXT_TOKEN_EDGE_NAME = 'CONTEXT_TOKEN'
# NATURALIZE_FEATURE_NODE_NODE_NAME = 'FEATURE_NODE'
# NATURALIZE_NEXT_FEATURE_EDGE_NAME = 'NEXT_FEATURE'
# NATURALIZE_FEATURE_ROOT_NODE_NAME = 'FEATURE_ROOT'
# NATURALIZE_FEATURE_EDGE_NAME = 'FEATURE'
# NATURALIZE_NATURALIZE_ROOT_NODE_NAME = 'NATURALIZE_ROOT'
# NATURALIZE_LOCAL_CONTEXT_EDGE_NAME = 'LOCAL_CONTEXT'
# NATURALIZE_FEATURES_EDGE_NAME = 'FEATURES'


import json
import os
import shutil
import apache_beam as beam
import git
from plur.utils import constants
  wild.

  Args:
    s: A string.

  Returns:
    A string where all the special characters are escaped.
  """
  # We use [1:-1] to remove the quote added by repr()
  return repr(s)[1:-1]


def graph_to_output_example_is_not_none(graph_to_output_example_dict):
  return graph_to_output_example_dict['GraphToOutputExample'] is not None


def filter_split(data, split):
  return data['split'] == split


def add_field_to_metadata(metadata, fieldname, value, distinguishing_text=''):
  metadata = (
      metadata
      | f'Add {fieldname} to metadata {distinguishing_text}' >> beam.ParDo(
          AddDictField(fieldname), value))
  return metadata


def filter_impossible_tfexample(tfexample_feature_dict,
                                output_token_vocab_dict):

# Path: plur/utils/constants.py
# STAGE_1_DIRNAME = 'stage_1'
# RAW_DATA_DIRNAME = 'raw_data'
# GIT_REPO_DIRNAME = 'git_repo'
# GRAPH_TO_OUTPUT_EXAMPLE_DIRNAME = 'graph_to_output_example'
# TRAIN_SPLIT_NAME = 'train'
# VALIDATION_SPLIT_NAME = 'valid'
# TEST_SPLIT_NAME = 'test'
# METADATA_FILENAME = 'dataset_metadata.json'
# STAGE_2_DIRNAME = 'stage_2'
# VOCAB_FILES_DIRNAME = 'vocab_files'
# NODE_TYPE_VOCAB_FILENAME = 'node_type_vocab.txt'
# NODE_LABEL_VOCAB_FILENAME = 'node_label_vocab.txt'
# EDGE_TYPE_VOCAB_FILENAME = 'edge_type_vocab.txt'
# OUTPUT_TOKEN_VOCAB_FILENAME = 'output_token_vocab.txt'
# TFRECORD_DIRNAME = 'tfrecords'
# POINTER_TOKEN = 'POINTER'
# CLASS_TOKEN_PREFIX = 'CLASS_'
# DONE_TOKEN = 'DONE'
# PAD_TOKEN = '<pad>'
# OOV_TOKEN = '<OOV>'
# RESERVED_TOKENS = [OOV_TOKEN, PAD_TOKEN]
# NUM_RESERVED_TOKENS = len(RESERVED_TOKENS)
# EXAMPLE_GUID = 'graph/global_id'  # For future use.
# KEY_PROVENANCE = 'metadata/provenance'
# KEY_OUTPUT_TOKEN_IDS = 'targets/token_ids'
# KEY_OUTPUT_TOKEN_TEXTS = 'targets/texts'
# KEY_COPY_INPUT_INDICES = 'nodes/copy_indices'
# KEY_COPY_OUTPUT_INDICES = 'targets/copy_indices'
# KEY_POINTER_INPUT_INDICES = 'nodes/pointer_indices'
# KEY_POINTER_OUTPUT_INDICES = 'targets/pointer_indices'
# KEY_NODE_TOKEN_IDS = 'nodes/node_id_sequences'
# KEY_NODE_TEXTS = 'nodes/node_text_sequences'
# KEY_NODE_TOKEN_LENGTHS = 'nodes/node_sequence_lengths'
# KEY_NODE_GUIDS = 'nodes/global_node_id'  # For future use.
# KEY_NODE_TYPE_IDS = 'nodes/node_types'
# KEY_EDGE_TYPE_IDS = 'edges/types'
# KEY_EDGE_SOURCE_INDICES = 'edges/source_indices'
# KEY_EDGE_DESTINATION_INDICES = 'edges/dest_indices'
# KEY_MASKING_CANDIDATE_INDICES = 'masking/candidate_indices'
# PROVENANCE_TENSOR_NAME = 'provenance'
# TARGET_TOKEN_IDS_TENSOR_NAME = 'target_token_ids'
# COPY_INDICES_TENSOR_NAME = 'copy_indices'
# POINTER_INDICES_TENSOR_NAME = 'pointer_indices'
# NODE_ID_SEQUENCES_TENSOR_NAME = 'node_id_sequences'
# NODE_TYPES_TENSOR_NAME = 'node_types'
# EDGE_INDICATORS_TENSOR_NAME = 'edge_indicators'
# NODE_TEXT_SEQUENCES_TENSOR_NAME = 'node_text_sequences'
# TARGET_TEXTS_TENSOR_NAME = 'target_texts'
# MASKING_CANDIDATE_TENSOR_NAME = 'candidate_indices'
# HOPPITY_ADD_NODE_OP_NAME = 'add_node'
# HOPPITY_DEL_NODE_OP_NAME = 'del_node'
# HOPPITY_REPLACE_VAL_OP_NAME = 'replace_val'
# HOPPITY_REPLACE_TYPE_OP_NAME = 'replace_type'
# HOPPITY_REPLACE_NOOP_OP_NAME = 'NoOp'
# HOPPITY_OUTPUT_UNKNOWN_TOKEN = 'UNKNOWN'
# NATURALIZE_CONTEXT_TOKEN_NODE_NAME = 'CONTEXT_TOKEN'
# NATURALIZE_NEXT_TOKEN_EDGE_NAME = 'NEXT_TOKEN'
# NATURALIZE_CONTEXT_ROOT_NODE_NAME = 'CONTEXT_ROOT'
# NATURALIZE_CONTEXT_TOKEN_EDGE_NAME = 'CONTEXT_TOKEN'
# NATURALIZE_FEATURE_NODE_NODE_NAME = 'FEATURE_NODE'
# NATURALIZE_NEXT_FEATURE_EDGE_NAME = 'NEXT_FEATURE'
# NATURALIZE_FEATURE_ROOT_NODE_NAME = 'FEATURE_ROOT'
# NATURALIZE_FEATURE_EDGE_NAME = 'FEATURE'
# NATURALIZE_NATURALIZE_ROOT_NODE_NAME = 'NATURALIZE_ROOT'
# NATURALIZE_LOCAL_CONTEXT_EDGE_NAME = 'LOCAL_CONTEXT'
# NATURALIZE_FEATURES_EDGE_NAME = 'FEATURES'


# Path: plur/utils/util.py
import json
import os
import shutil
import apache_beam as beam
import git
from plur.utils import constants

  wild.

  Args:
    s: A string.

  Returns:
    A string where all the special characters are escaped.
  """
  # We use [1:-1] to remove the quote added by repr()
  return repr(s)[1:-1]


def graph_to_output_example_is_not_none(graph_to_output_example_dict):
  return graph_to_output_example_dict['GraphToOutputExample'] is not None


def filter_split(data, split):
  return data['split'] == split


def add_field_to_metadata(metadata, fieldname, value, distinguishing_text=''):
  metadata = (
      metadata
      | f'Add {fieldname} to metadata {distinguishing_text}' >> beam.ParDo(
          AddDictField(fieldname), value))
  return metadata


def filter_impossible_tfexample(tfexample_feature_dict,
                                output_token_vocab_dict):

  output_oov_token_id = output_token_vocab_dict[constants.OOV_TOKEN]