From 5efb51d566b0da70f2aabda9bb6648695af85756 Mon Sep 17 00:00:00 2001
From: zoy <zoy@google.com>
Date: Thu, 4 Aug 2022 11:54:03 -0700
Subject: [PATCH] Cleanup: add exception chaining to graph_tools errors.

PiperOrigin-RevId: 465369075
---
 tensorflow_transform/graph_tools.py      | 26 ++++++++++--------------
 tensorflow_transform/graph_tools_test.py |  4 +---
 2 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/tensorflow_transform/graph_tools.py b/tensorflow_transform/graph_tools.py
index 80d90220..da91d9bb 100644
--- a/tensorflow_transform/graph_tools.py
+++ b/tensorflow_transform/graph_tools.py
@@ -179,31 +179,27 @@ def wrapper(self, tensor_or_op):
     try:
       return func(self, tensor_or_op)
     except _UnexpectedPlaceholderError as e:
-      if e.func_graph_name:
-        raise ValueError(
-            'The tensor_or_op {} depended on a placeholder ({}) that is part '
-            'of a tf.function graph ({}), this is not supported. This may be a '
-            'result of calling a tf.Transform analyzer in a tf.function'
-            ''.format(tensor_or_op, e.tensor, e.func_graph_name))
-      else:
-        raise ValueError(
-            'The tensor_or_op {} depended on a placeholder ({}) that was not '
-            'in the input_signature.  This may have be caused by manually '
-            'adding a placeholder to the graph'.format(tensor_or_op, e.tensor))
+      context = (f' tf.function name: `{e.func_graph_name}`'
+                 if e.func_graph_name else '')
+      raise ValueError(
+          'The tensor_or_op {} depended on a placeholder ({}) that was not '
+          'in the input_signature.  This may have be caused by manually '
+          'adding a placeholder to the graph.{}'.format(tensor_or_op, e.tensor,
+                                                        context)) from e
     except _UnexpectedTableError as e:
       if e.func_graph_name:
         raise ValueError(
             'The tensor_or_op {} depended on an initializable table ({}) that '
             'is part of a tf.function graph ({}), this is not supported. This'
             ' may be a result of initializing a table in a tf.function'
-            ''.format(tensor_or_op, e.op, e.func_graph_name))
+            ''.format(tensor_or_op, e.op, e.func_graph_name)) from e
       else:
         raise ValueError(
             'The tensor_or_op {} depended on an initializable table ({}) that '
             'was not tracked by the graph analysis.  This may be caused by '
             'adding an initializable table without adding its initializer to '
             'the collection tf.GraphKeys.TABLE_INITIALIZERS'.format(
-                tensor_or_op, e.op))
+                tensor_or_op, e.op)) from e
 
   return wrapper
 
@@ -687,7 +683,7 @@ def _get_table_init_op_source_info(self, table_init_op, graph_analyzer,
       raise ValueError(
           'The table initializer {} depended on a placeholder ({}).  Note '
           'placeholders will not be fed during table initialization'.format(
-              table_init_op, e.tensor))
+              table_init_op, e.tensor)) from e
     except _UnexpectedTableError as e:
       if e.func_graph_name:
         raise e
@@ -695,7 +691,7 @@ def _get_table_init_op_source_info(self, table_init_op, graph_analyzer,
           'The table initializer {} depended on an initializable table ({}). '
           'Note tables are initialized in one pass so a table initializer '
           'cannot depend on the output of an initializeable table'.format(
-              table_init_op, e.op))
+              table_init_op, e.op)) from e
     return _SourceInfo(ready, path)
 
   @property
diff --git a/tensorflow_transform/graph_tools_test.py b/tensorflow_transform/graph_tools_test.py
index 3ea062ef..27a195d6 100644
--- a/tensorflow_transform/graph_tools_test.py
+++ b/tensorflow_transform/graph_tools_test.py
@@ -636,9 +636,7 @@ def testInitializableGraphAnalyzerConstructorRaises(
       (_create_graph_with_y_function_of_x, [], {}, 'y',
        'may have be caused by manually adding a placeholder to the graph'),
       (_create_graph_with_placeholder_in_tf_function, ['x'], {}, 'z',
-       r'that is part of a tf.function graph \(foo\), this is not supported. '  # pylint: disable=implicit-str-concat
-       'This may be a result of calling a tf.Transform analyzer in a '
-       'tf.function'),
+       'manually adding a placeholder to the graph. tf.function name: `foo`'),
       (_create_graph_with_y_function_of_x_and_untracked_table, ['x'], {
           'filename': True
       }, 'y', 'may be caused by adding an initializable table without'),