Handle newly added lifecyle_stage.DISABLED

PiperOrigin-RevId: 345814659

Author: tfx-copybara

struct2tensor/expression_impl/promote.py

@@ -213,10 +213,11 @@ def _lifecycle_stage_number(a) -> int:
     An integer that corresponds to the lifecycle stage of 'a'.
   """
   stages = [
-      schema_pb2.LifecycleStage.DEPRECATED, schema_pb2.LifecycleStage.PLANNED,
-      schema_pb2.LifecycleStage.ALPHA, schema_pb2.LifecycleStage.DEBUG_ONLY,
-      None, schema_pb2.LifecycleStage.UNKNOWN_STAGE,
-      schema_pb2.LifecycleStage.BETA, schema_pb2.LifecycleStage.PRODUCTION
+      schema_pb2.LifecycleStage.DEPRECATED, schema_pb2.LifecycleStage.DISABLED,
+      schema_pb2.LifecycleStage.PLANNED, schema_pb2.LifecycleStage.ALPHA,
+      schema_pb2.LifecycleStage.DEBUG_ONLY, None,
+      schema_pb2.LifecycleStage.UNKNOWN_STAGE, schema_pb2.LifecycleStage.BETA,
+      schema_pb2.LifecycleStage.PRODUCTION
   ]
   return stages.index(a)
 

struct2tensor/expression_impl/promote_test.py

@@ -106,6 +106,7 @@ def test_promote_with_schema_dense_parent(self):
   def test_lifecycle_stage(self):
     # Stages have the following priority, from lowest to highest:
     #   schema_pb2.LifecycleStage.DEPRECATED
+    #   schema_pb2.LifecycleStage.DISABLED
     #   schema_pb2.LifecycleStage.PLANNED,
     #   schema_pb2.LifecycleStage.ALPHA
     #   schema_pb2.LifecycleStage.DEBUG_ONLY,
@@ -139,7 +140,17 @@ def _check_lifecycle_stage(a, b):
     self.assertEqual(
         schema_pb2.LifecycleStage.DEPRECATED,
         _check_lifecycle_stage(schema_pb2.LifecycleStage.DEPRECATED,
-                               schema_pb2.LifecycleStage.PLANNED))
+                               schema_pb2.LifecycleStage.DISABLED))
+    self.assertEqual(
+        schema_pb2.LifecycleStage.DEPRECATED,
+        _check_lifecycle_stage(schema_pb2.LifecycleStage.DISABLED,
+                               schema_pb2.LifecycleStage.DEPRECATED))
+
+    self.assertEqual(
+        schema_pb2.LifecycleStage.DISABLED,
+        _check_lifecycle_stage(schema_pb2.LifecycleStage.PLANNED,
+                               schema_pb2.LifecycleStage.DISABLED))
+
     self.assertEqual(
         schema_pb2.LifecycleStage.DEPRECATED,
         _check_lifecycle_stage(schema_pb2.LifecycleStage.PLANNED,