Add a helper function to create an empty placeholder context and upda…

…te tests to use the helper. PiperOrigin-RevId: 650300708
tensorflow · Jul 9, 2024 · e1fd430 · e1fd430
1 parent fffb9d7
commit e1fd430
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/tfx/dsl/compiler/placeholder_utils.py b/tfx/dsl/compiler/placeholder_utils.py
@@ -129,6 +129,13 @@ def resolve_placeholder_expression(
   return result
 
 
+def empty_placeholder_context() -> ResolutionContext:
+  """Returns an empty placeholder context."""
+  return ResolutionContext(
+      exec_info=data_types.ExecutionInfo(),
+  )
+
+
 class _Operation(enum.Enum):
   """Alias for Operation enum types in placeholder.proto."""
 

diff --git a/tfx/dsl/compiler/placeholder_utils_test.py b/tfx/dsl/compiler/placeholder_utils_test.py
@@ -1689,8 +1689,8 @@ def testMakeProtoOpResolvesProto(self):
         placeholder_pb2.PlaceholderExpression(),
     )
     resolved_proto = placeholder_utils.resolve_placeholder_expression(
-        placeholder_expression, placeholder_utils.ResolutionContext(
-            exec_info=data_types.ExecutionInfo()))
+        placeholder_expression, placeholder_utils.empty_placeholder_context()
+    )
     self.assertProtoEquals(
         """
         splits: "train"

diff --git a/tfx/dsl/placeholder/placeholder_test_util.py b/tfx/dsl/placeholder/placeholder_test_util.py
@@ -17,7 +17,6 @@
 
 from tfx.dsl.compiler import placeholder_utils
 from tfx.dsl.placeholder import placeholder_base
-from tfx.orchestration.portable import data_types
 
 
 def resolve(
@@ -39,9 +38,7 @@ def resolve(
   return placeholder_utils.resolve_placeholder_expression(
       placeholder.encode(),
       resolution_context
-      or placeholder_utils.ResolutionContext(
-          exec_info=data_types.ExecutionInfo()
-      ),
+      or placeholder_utils.empty_placeholder_context(),
   )