# Path: tensorflow_transform/annotators.py
# _ASSET_KEY_COLLECTION = 'tft_asset_key_collection'
# _ASSET_FILENAME_COLLECTION = 'tft_asset_filename_collection'
# _OBJECT_TRACKER = None
# VOCABULARY_SIZE_BY_NAME_COLLECTION = 'tft_vocabulary_size_by_name_collection'
#   _OBJECT_TRACKER = object_tracker
#     _OBJECT_TRACKER = None
# class ObjectTracker:
#   def __init__(self):
#   def trackable_objects(self) -> List[base.Trackable]:
#   def add_trackable_object(self, trackable_object: base.Trackable,
#                            name: Optional[str]):
# def object_tracker_scope(object_tracker: ObjectTracker):
# def _get_object(name: str) -> Optional[base.Trackable]:
# def track_object(trackable: base.Trackable, name: Optional[str]):
# def make_and_track_object(trackable_factory_callable: Callable[[],
#                                                                base.Trackable],
#                           name: Optional[str] = None) -> base.Trackable:
# def get_asset_annotations(graph: tf.Graph):
# def clear_asset_annotations(graph: tf.Graph):
# def annotate_asset(asset_key: str, asset_filename: str):
# def annotate_vocab_size(vocab_filename: str, vocab_size: tf.Tensor):
# 
# Path: tensorflow_transform/test_case.py
# def is_tf_api_version_1():
# def cross_named_parameters(*args):
#   def _cross_test_cases(parameters_list):
# def parameters(*testcases):
#   def wrapper(fn):
#     def to_arg_dict(testcase):
# def cross_parameters(*args):
# def _make_placeholder(tensor_spec):
# def _graph_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_graph(*inputs):
# def _ragged_value_as_constant(value, dtype):
# def _wrap_as_constant(value, tensor_spec):
# def _eager_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_eagerly(*inputs):  # pylint: disable=missing-docstring
# def _tf_function_function_handler(input_signature):
#   def wrapper(fn):
# def is_external_environment():
# def skip_if_external_environment(reason):
# def skip_if_not_tf2(reason):
# def cross_with_function_handlers(parameters_list):
#   def assertDataCloseOrEqual(self, a_data, b_data):
#   def _assertValuesCloseOrEqual(self, a_value, b_value, msg=None):
#   def AssertVocabularyContents(self, vocab_file_path, file_contents):
#   def WriteRenderedDotFile(self, dot_string, output_file=None):
#   def _NumpyArraysToLists(self, maybe_arrays):
#   def _SortedDicts(self, list_of_dicts):
#   def _SortedData(self, list_of_dicts_of_arrays):
# FUNCTION_HANDLERS = [
#     dict(testcase_name='graph',
#          function_handler=_graph_function_handler),
#     dict(testcase_name='eager',
#          function_handler=_eager_function_handler),
#     dict(testcase_name='tf_function',
#          function_handler=_tf_function_function_handler)
# ]
# class TransformTestCase(parameterized.TestCase, tf.test.TestCase):


import tensorflow as tf
from tensorflow_transform import annotators
from tensorflow_transform import test_case
from tensorflow.python.training.tracking import base  # pylint: disable=g-direct-tensorflow-import
      graph = tf.function(foo).get_concrete_function().graph

    self.assertDictEqual(
        annotators.get_asset_annotations(graph), {
            'my_key': 'my_value',
            'my_key2': 'my_value2'
        })

    annotators.clear_asset_annotations(graph)
    self.assertDictEqual(annotators.get_asset_annotations(graph), {})

  def test_object_tracker(self):
    test_case.skip_if_not_tf2('Tensorflow 2.x required')

    trackable_object = base.Trackable()

    @tf.function
    def preprocessing_fn():
      _ = annotators.make_and_track_object(lambda: trackable_object)
      return 1

    object_tracker = annotators.ObjectTracker()
    with annotators.object_tracker_scope(object_tracker):
      _ = preprocessing_fn()

    self.assertLen(object_tracker.trackable_objects, 1)
    self.assertEqual(trackable_object, object_tracker.trackable_objects[0])


if __name__ == '__main__':

# Path: tensorflow_transform/annotators.py
# _ASSET_KEY_COLLECTION = 'tft_asset_key_collection'
# _ASSET_FILENAME_COLLECTION = 'tft_asset_filename_collection'
# _OBJECT_TRACKER = None
# VOCABULARY_SIZE_BY_NAME_COLLECTION = 'tft_vocabulary_size_by_name_collection'
#   _OBJECT_TRACKER = object_tracker
#     _OBJECT_TRACKER = None
# class ObjectTracker:
#   def __init__(self):
#   def trackable_objects(self) -> List[base.Trackable]:
#   def add_trackable_object(self, trackable_object: base.Trackable,
#                            name: Optional[str]):
# def object_tracker_scope(object_tracker: ObjectTracker):
# def _get_object(name: str) -> Optional[base.Trackable]:
# def track_object(trackable: base.Trackable, name: Optional[str]):
# def make_and_track_object(trackable_factory_callable: Callable[[],
#                                                                base.Trackable],
#                           name: Optional[str] = None) -> base.Trackable:
# def get_asset_annotations(graph: tf.Graph):
# def clear_asset_annotations(graph: tf.Graph):
# def annotate_asset(asset_key: str, asset_filename: str):
# def annotate_vocab_size(vocab_filename: str, vocab_size: tf.Tensor):
# 
# Path: tensorflow_transform/test_case.py
# def is_tf_api_version_1():
# def cross_named_parameters(*args):
#   def _cross_test_cases(parameters_list):
# def parameters(*testcases):
#   def wrapper(fn):
#     def to_arg_dict(testcase):
# def cross_parameters(*args):
# def _make_placeholder(tensor_spec):
# def _graph_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_graph(*inputs):
# def _ragged_value_as_constant(value, dtype):
# def _wrap_as_constant(value, tensor_spec):
# def _eager_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_eagerly(*inputs):  # pylint: disable=missing-docstring
# def _tf_function_function_handler(input_signature):
#   def wrapper(fn):
# def is_external_environment():
# def skip_if_external_environment(reason):
# def skip_if_not_tf2(reason):
# def cross_with_function_handlers(parameters_list):
#   def assertDataCloseOrEqual(self, a_data, b_data):
#   def _assertValuesCloseOrEqual(self, a_value, b_value, msg=None):
#   def AssertVocabularyContents(self, vocab_file_path, file_contents):
#   def WriteRenderedDotFile(self, dot_string, output_file=None):
#   def _NumpyArraysToLists(self, maybe_arrays):
#   def _SortedDicts(self, list_of_dicts):
#   def _SortedData(self, list_of_dicts_of_arrays):
# FUNCTION_HANDLERS = [
#     dict(testcase_name='graph',
#          function_handler=_graph_function_handler),
#     dict(testcase_name='eager',
#          function_handler=_eager_function_handler),
#     dict(testcase_name='tf_function',
#          function_handler=_tf_function_function_handler)
# ]
# class TransformTestCase(parameterized.TestCase, tf.test.TestCase):


# Path: tensorflow_transform/annotators_test.py
import tensorflow as tf
from tensorflow_transform import annotators
from tensorflow_transform import test_case
from tensorflow.python.training.tracking import base  # pylint: disable=g-direct-tensorflow-import

      graph = tf.function(foo).get_concrete_function().graph

    self.assertDictEqual(
        annotators.get_asset_annotations(graph), {
            'my_key': 'my_value',
            'my_key2': 'my_value2'
        })

    annotators.clear_asset_annotations(graph)
    self.assertDictEqual(annotators.get_asset_annotations(graph), {})

  def test_object_tracker(self):
    test_case.skip_if_not_tf2('Tensorflow 2.x required')

    trackable_object = base.Trackable()

    @tf.function
    def preprocessing_fn():
      _ = annotators.make_and_track_object(lambda: trackable_object)
      return 1

    object_tracker = annotators.ObjectTracker()
    with annotators.object_tracker_scope(object_tracker):
      _ = preprocessing_fn()

    self.assertLen(object_tracker.trackable_objects, 1)
    self.assertEqual(trackable_object, object_tracker.trackable_objects[0])


if __name__ == '__main__':

  test_case.main()