# Path: tensorflow_transform/beam/tft_unit.py
# def metadata_from_feature_spec(feature_spec, domains=None):
# def canonical_numeric_dtype(dtype):
# def make_feature_spec_wrapper(make_feature_spec, *args):
# def _format_example_as_numpy_dict(example, feature_shape_dict):
#     def has_ran(self):
#     def metrics(self):
#     def __exit__(self, exc_type, exc_val, exc_tb):
#   def _makeTestPipeline(self):
#   def assertMetricsCounterEqual(self, metrics, name, expected_count,
#                                 namespaces_list=None):
#   def assertAnalyzerOutputs(self,
#                             input_data,
#                             input_metadata,
#                             analyzer_fn,
#                             expected_outputs,
#                             test_data=None,
#                             desired_batch_size=None,
#                             beam_pipeline=None,
#                             force_tf_compat_v1=False,
#                             output_record_batches=False):
#     def preprocessing_fn(inputs):
#   def assertAnalyzeAndTransformResults(self,
#                                        input_data,
#                                        input_metadata,
#                                        preprocessing_fn,
#                                        expected_data=None,
#                                        expected_metadata=None,
#                                        expected_vocab_file_contents=None,
#                                        test_data=None,
#                                        desired_batch_size=None,
#                                        beam_pipeline=None,
#                                        temp_dir=None,
#                                        force_tf_compat_v1=False,
#                                        output_record_batches=False):
# class TransformTestCase(test_case.TransformTestCase):
#   class _TestPipeline(beam.Pipeline):
# 
# Path: tensorflow_transform/beam/tft_beam_io/beam_metadata_io.py
# class BeamDatasetMetadata(
#     tfx_namedtuple.namedtuple(
#         'BeamDatasetMetadata',
#         ['dataset_metadata', 'deferred_metadata', 'asset_map'])):
# class WriteMetadata(beam.PTransform):
#   def schema(self):
#   def __init__(self, path, pipeline, write_to_unique_subdirectory=False):
#   def _extract_input_pvalues(self, metadata):
#   def expand(self, metadata):
#     def write_metadata_output(metadata):
# 
# Path: tensorflow_transform/beam/tft_beam_io/transform_fn_io.py
# TRANSFORMED_METADATA_DIR = tft.TFTransformOutput.TRANSFORMED_METADATA_DIR
# TRANSFORM_FN_DIR = tft.TFTransformOutput.TRANSFORM_FN_DIR
# def _copy_tree_to_unique_temp_dir(source, base_temp_dir_path):
# def _copy_tree(source, destination):
#   def __init__(self, path):
#   def _extract_input_pvalues(self, transform_fn):
#   def expand(self, transform_fn):
#     def publish_outputs(unused_element, metadata_source_path,
#                         transform_fn_source_path):
#   def __init__(self, path):
#   def expand(self, pvalue):
# class WriteTransformFn(beam.PTransform):
# class ReadTransformFn(beam.PTransform):
# 
# Path: tensorflow_transform/beam/tft_beam_io/test_metadata.py
# _FEATURE_SPEC = {
#     'fixed_column': tf.io.FixedLenFeature([3], tf.string),
#     'list_columm': tf.io.VarLenFeature(tf.int64),
# }
# COMPLETE_METADATA = dataset_metadata.DatasetMetadata.from_feature_spec(
#     _FEATURE_SPEC, domains={'list_columm': schema_pb2.IntDomain(min=-1, max=5)})
# INCOMPLETE_METADATA = dataset_metadata.DatasetMetadata.from_feature_spec(
#     _FEATURE_SPEC,
#     # Values will be overridden by those in COMPLETE_METADATA
#     domains={'list_columm': schema_pb2.IntDomain(min=0, max=0)})
# 
# Path: tensorflow_transform/tf_metadata/metadata_io.py
# def read_metadata(path):
# def _parse_schema_json(schema_json):
# def _column_schema_from_json(feature_dict):
# def _domain_from_json(domain):
#     def maybe_to_int(s):
# def _dtype_from_json(domain):
# def write_metadata(metadata, path):
# def _convert_scalar_or_list(fn, scalar_or_list):


import os
import apache_beam as beam
import tensorflow as tf
import tensorflow_transform as tft
from apache_beam.testing import util as beam_test_util
from tensorflow_transform.beam import tft_unit
from tensorflow_transform.beam.tft_beam_io import beam_metadata_io
from tensorflow_transform.beam.tft_beam_io import transform_fn_io
from tensorflow_transform.beam.tft_beam_io import test_metadata
from tensorflow_transform.tf_metadata import metadata_io
from tensorflow.python.lib.io import file_io  # pylint: disable=g-direct-tensorflow-import
    # NOTE: we don't need to create or write to the transform_fn directory since
    # ReadTransformFn never inspects this directory.
    transform_fn_dir = os.path.join(
        path, tft.TFTransformOutput.TRANSFORM_FN_DIR)
    transformed_metadata_dir = os.path.join(
        path, tft.TFTransformOutput.TRANSFORMED_METADATA_DIR)
    metadata_io.write_metadata(test_metadata.COMPLETE_METADATA,
                               transformed_metadata_dir)

    with beam.Pipeline() as pipeline:
      saved_model_dir_pcoll, metadata = (
          pipeline | transform_fn_io.ReadTransformFn(path))
      beam_test_util.assert_that(
          saved_model_dir_pcoll,
          beam_test_util.equal_to([transform_fn_dir]),
          label='AssertSavedModelDir')
      # NOTE: metadata is currently read in a non-deferred manner.
      self.assertEqual(metadata, test_metadata.COMPLETE_METADATA)

  def testWriteTransformFn(self):
    transform_output_dir = os.path.join(self.get_temp_dir(), 'output')

    with beam.Pipeline() as pipeline:
      # Create an empty directory for the source saved model dir.
      saved_model_dir = os.path.join(self.get_temp_dir(), 'source')
      file_io.recursive_create_dir(saved_model_dir)
      saved_model_dir_pcoll = (
          pipeline | 'CreateSavedModelDir' >> beam.Create([saved_model_dir]))
      # Combine test metadata with a dict of PCollections resolving futures.
      deferred_metadata = pipeline | 'CreateDeferredMetadata' >> beam.Create(

# Path: tensorflow_transform/beam/tft_unit.py
# def metadata_from_feature_spec(feature_spec, domains=None):
# def canonical_numeric_dtype(dtype):
# def make_feature_spec_wrapper(make_feature_spec, *args):
# def _format_example_as_numpy_dict(example, feature_shape_dict):
#     def has_ran(self):
#     def metrics(self):
#     def __exit__(self, exc_type, exc_val, exc_tb):
#   def _makeTestPipeline(self):
#   def assertMetricsCounterEqual(self, metrics, name, expected_count,
#                                 namespaces_list=None):
#   def assertAnalyzerOutputs(self,
#                             input_data,
#                             input_metadata,
#                             analyzer_fn,
#                             expected_outputs,
#                             test_data=None,
#                             desired_batch_size=None,
#                             beam_pipeline=None,
#                             force_tf_compat_v1=False,
#                             output_record_batches=False):
#     def preprocessing_fn(inputs):
#   def assertAnalyzeAndTransformResults(self,
#                                        input_data,
#                                        input_metadata,
#                                        preprocessing_fn,
#                                        expected_data=None,
#                                        expected_metadata=None,
#                                        expected_vocab_file_contents=None,
#                                        test_data=None,
#                                        desired_batch_size=None,
#                                        beam_pipeline=None,
#                                        temp_dir=None,
#                                        force_tf_compat_v1=False,
#                                        output_record_batches=False):
# class TransformTestCase(test_case.TransformTestCase):
#   class _TestPipeline(beam.Pipeline):
# 
# Path: tensorflow_transform/beam/tft_beam_io/beam_metadata_io.py
# class BeamDatasetMetadata(
#     tfx_namedtuple.namedtuple(
#         'BeamDatasetMetadata',
#         ['dataset_metadata', 'deferred_metadata', 'asset_map'])):
# class WriteMetadata(beam.PTransform):
#   def schema(self):
#   def __init__(self, path, pipeline, write_to_unique_subdirectory=False):
#   def _extract_input_pvalues(self, metadata):
#   def expand(self, metadata):
#     def write_metadata_output(metadata):
# 
# Path: tensorflow_transform/beam/tft_beam_io/transform_fn_io.py
# TRANSFORMED_METADATA_DIR = tft.TFTransformOutput.TRANSFORMED_METADATA_DIR
# TRANSFORM_FN_DIR = tft.TFTransformOutput.TRANSFORM_FN_DIR
# def _copy_tree_to_unique_temp_dir(source, base_temp_dir_path):
# def _copy_tree(source, destination):
#   def __init__(self, path):
#   def _extract_input_pvalues(self, transform_fn):
#   def expand(self, transform_fn):
#     def publish_outputs(unused_element, metadata_source_path,
#                         transform_fn_source_path):
#   def __init__(self, path):
#   def expand(self, pvalue):
# class WriteTransformFn(beam.PTransform):
# class ReadTransformFn(beam.PTransform):
# 
# Path: tensorflow_transform/beam/tft_beam_io/test_metadata.py
# _FEATURE_SPEC = {
#     'fixed_column': tf.io.FixedLenFeature([3], tf.string),
#     'list_columm': tf.io.VarLenFeature(tf.int64),
# }
# COMPLETE_METADATA = dataset_metadata.DatasetMetadata.from_feature_spec(
#     _FEATURE_SPEC, domains={'list_columm': schema_pb2.IntDomain(min=-1, max=5)})
# INCOMPLETE_METADATA = dataset_metadata.DatasetMetadata.from_feature_spec(
#     _FEATURE_SPEC,
#     # Values will be overridden by those in COMPLETE_METADATA
#     domains={'list_columm': schema_pb2.IntDomain(min=0, max=0)})
# 
# Path: tensorflow_transform/tf_metadata/metadata_io.py
# def read_metadata(path):
# def _parse_schema_json(schema_json):
# def _column_schema_from_json(feature_dict):
# def _domain_from_json(domain):
#     def maybe_to_int(s):
# def _dtype_from_json(domain):
# def write_metadata(metadata, path):
# def _convert_scalar_or_list(fn, scalar_or_list):


# Path: tensorflow_transform/beam/tft_beam_io/transform_fn_io_test.py
import os
import apache_beam as beam
import tensorflow as tf
import tensorflow_transform as tft
from apache_beam.testing import util as beam_test_util
from tensorflow_transform.beam import tft_unit
from tensorflow_transform.beam.tft_beam_io import beam_metadata_io
from tensorflow_transform.beam.tft_beam_io import transform_fn_io
from tensorflow_transform.beam.tft_beam_io import test_metadata
from tensorflow_transform.tf_metadata import metadata_io
from tensorflow.python.lib.io import file_io  # pylint: disable=g-direct-tensorflow-import

    # NOTE: we don't need to create or write to the transform_fn directory since
    # ReadTransformFn never inspects this directory.
    transform_fn_dir = os.path.join(
        path, tft.TFTransformOutput.TRANSFORM_FN_DIR)
    transformed_metadata_dir = os.path.join(
        path, tft.TFTransformOutput.TRANSFORMED_METADATA_DIR)
    metadata_io.write_metadata(test_metadata.COMPLETE_METADATA,
                               transformed_metadata_dir)

    with beam.Pipeline() as pipeline:
      saved_model_dir_pcoll, metadata = (
          pipeline | transform_fn_io.ReadTransformFn(path))
      beam_test_util.assert_that(
          saved_model_dir_pcoll,
          beam_test_util.equal_to([transform_fn_dir]),
          label='AssertSavedModelDir')
      # NOTE: metadata is currently read in a non-deferred manner.
      self.assertEqual(metadata, test_metadata.COMPLETE_METADATA)

  def testWriteTransformFn(self):
    transform_output_dir = os.path.join(self.get_temp_dir(), 'output')

    with beam.Pipeline() as pipeline:
      # Create an empty directory for the source saved model dir.
      saved_model_dir = os.path.join(self.get_temp_dir(), 'source')
      file_io.recursive_create_dir(saved_model_dir)
      saved_model_dir_pcoll = (
          pipeline | 'CreateSavedModelDir' >> beam.Create([saved_model_dir]))
      # Combine test metadata with a dict of PCollections resolving futures.
      deferred_metadata = pipeline | 'CreateDeferredMetadata' >> beam.Create(

          [test_metadata.COMPLETE_METADATA])