# Path: tensorflow_transform/nodes.py
# class ValueNode(
#     tfx_namedtuple.namedtuple('ValueNode',
#                               ['parent_operation', 'value_index'])):
# class OperationDef(metaclass=abc.ABCMeta):
# class OperationNode:
# class Visitor(metaclass=abc.ABCMeta):
# class Traverser:
# class _PrintGraphVisitor(Visitor):
#   def __init__(self, parent_operation, value_index: int):
#   def __iter__(self):
#   def num_outputs(self) -> int:
#   def label(self) -> str:
#   def get_field_str(self, field_name: str) -> str:
#   def is_partitionable(self) -> bool:
#   def cache_coder(self) -> Optional[object]:
#   def __init__(self, operation_def, inputs):
#   def __repr__(self):
#   def operation_def(self):
#   def inputs(self):
#   def outputs(self):
# def apply_operation(operation_def_cls, *args, **kwargs):
# def apply_multi_output_operation(operation_def_cls, *args, **kwargs):
#   def validate_value(self, value):
#   def visit(self, operation_def, input_values):
#   def __init__(self, visitor: Visitor):
#   def visit_value_node(self, value_node: ValueNode):
#   def _maybe_visit_value_node(self, value_node: ValueNode):
#   def _visit_operation(self, operation: OperationNode):
# def _escape(line: str) -> str:
#   def __init__(self):
#   def get_dot_graph(self) -> pydot.Dot:
#   def visit(self, operation_def, input_nodes) -> Tuple[pydot.Node, ...]:
#   def validate_value(self, value: pydot.Node):
# def get_dot_graph(leaf_nodes: Collection[ValueNode]) -> pydot.Dot:
# 
# 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 nodes
from tensorflow_transform import test_case
from tfx_bsl.types import tfx_namedtuple
  def testGetDotGraph(self):
    a = nodes.apply_operation(_Constant, value='a', label='Constant[a]')
    b = nodes.apply_operation(_Constant, value='b', label='Constant[b]')
    b_copy, a_copy = nodes.apply_multi_output_operation(
        _Swap, a, b, label='Swap[0]')
    b_copy2, unused_a_copy2 = nodes.apply_multi_output_operation(
        _Swap, a_copy, b_copy, label='Swap[1]')
    dot_string = nodes.get_dot_graph([b_copy2]).to_string()
    self.WriteRenderedDotFile(dot_string)

    self.assertMultiLineEqual(
        dot_string,
        """\
digraph G {
directed=True;
node [shape=Mrecord];
"Constant[a]" [label="{_Constant|value: a|label: Constant[a]}"];
"Constant[b]" [label="{_Constant|value: b|label: Constant[b]}"];
"Swap[0]" [label="{_Swap|label: Swap[0]|{<0>0|<1>1}}"];
"Constant[a]" -> "Swap[0]";
"Constant[b]" -> "Swap[0]";
"Swap[1]" [label="{_Swap|label: Swap[1]|{<0>0|<1>1}}"];
"Swap[0]":1 -> "Swap[1]";
"Swap[0]":0 -> "Swap[1]";
}
""",
        msg='Result dot graph is:\n{}'.format(dot_string))


if __name__ == '__main__':

# Path: tensorflow_transform/nodes.py
# class ValueNode(
#     tfx_namedtuple.namedtuple('ValueNode',
#                               ['parent_operation', 'value_index'])):
# class OperationDef(metaclass=abc.ABCMeta):
# class OperationNode:
# class Visitor(metaclass=abc.ABCMeta):
# class Traverser:
# class _PrintGraphVisitor(Visitor):
#   def __init__(self, parent_operation, value_index: int):
#   def __iter__(self):
#   def num_outputs(self) -> int:
#   def label(self) -> str:
#   def get_field_str(self, field_name: str) -> str:
#   def is_partitionable(self) -> bool:
#   def cache_coder(self) -> Optional[object]:
#   def __init__(self, operation_def, inputs):
#   def __repr__(self):
#   def operation_def(self):
#   def inputs(self):
#   def outputs(self):
# def apply_operation(operation_def_cls, *args, **kwargs):
# def apply_multi_output_operation(operation_def_cls, *args, **kwargs):
#   def validate_value(self, value):
#   def visit(self, operation_def, input_values):
#   def __init__(self, visitor: Visitor):
#   def visit_value_node(self, value_node: ValueNode):
#   def _maybe_visit_value_node(self, value_node: ValueNode):
#   def _visit_operation(self, operation: OperationNode):
# def _escape(line: str) -> str:
#   def __init__(self):
#   def get_dot_graph(self) -> pydot.Dot:
#   def visit(self, operation_def, input_nodes) -> Tuple[pydot.Node, ...]:
#   def validate_value(self, value: pydot.Node):
# def get_dot_graph(leaf_nodes: Collection[ValueNode]) -> pydot.Dot:
# 
# 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/nodes_test.py
import tensorflow as tf
from tensorflow_transform import nodes
from tensorflow_transform import test_case
from tfx_bsl.types import tfx_namedtuple

  def testGetDotGraph(self):
    a = nodes.apply_operation(_Constant, value='a', label='Constant[a]')
    b = nodes.apply_operation(_Constant, value='b', label='Constant[b]')
    b_copy, a_copy = nodes.apply_multi_output_operation(
        _Swap, a, b, label='Swap[0]')
    b_copy2, unused_a_copy2 = nodes.apply_multi_output_operation(
        _Swap, a_copy, b_copy, label='Swap[1]')
    dot_string = nodes.get_dot_graph([b_copy2]).to_string()
    self.WriteRenderedDotFile(dot_string)

    self.assertMultiLineEqual(
        dot_string,
        """\
digraph G {
directed=True;
node [shape=Mrecord];
"Constant[a]" [label="{_Constant|value: a|label: Constant[a]}"];
"Constant[b]" [label="{_Constant|value: b|label: Constant[b]}"];
"Swap[0]" [label="{_Swap|label: Swap[0]|{<0>0|<1>1}}"];
"Constant[a]" -> "Swap[0]";
"Constant[b]" -> "Swap[0]";
"Swap[1]" [label="{_Swap|label: Swap[1]|{<0>0|<1>1}}"];
"Swap[0]":1 -> "Swap[1]";
"Swap[0]":0 -> "Swap[1]";
}
""",
        msg='Result dot graph is:\n{}'.format(dot_string))


if __name__ == '__main__':

  test_case.main()