Add support for custom_value in GenC Python API.

PiperOrigin-RevId: 658890970

genc/python/base/BUILD

@@ -78,6 +78,8 @@ py_library(
     deps = [
         ":computation",
         "//genc/proto/v0:computation_py_pb2",
+        "//google/protobuf:any_py_pb2",
+        "//net/proto2/python/public",
     ],
 )
 
@@ -87,5 +89,6 @@ py_test(
     deps = [
         ":to_from_value_proto",
         "//genc/proto/v0:computation_py_pb2",
+        "//google/protobuf:timestamp_py_pb2",
     ],
 )

genc/python/base/to_from_value_proto.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 """Libraries fo packing and unpacking to/from `Value` proto."""
 
+from google3.google.protobuf import any_pb2
+from google3.net.proto2.python.public import message
 from genc.proto.v0 import computation_pb2 as pb
 from genc.python.base import computation
 
@@ -50,6 +52,14 @@ def to_value_proto(arg):
       element.label = key
       struct.element.append(element)
     return pb.Value(struct=struct)
+  elif isinstance(arg, message.Message):
+    # Create a new Any message
+    any_message = any_pb2.Any()
+    # Pack the input message into the Any message
+    any_message.Pack(arg)
+    # Create a Value message and set its any_val field
+    value_message = pb.Value(custom_value=any_message)
+    return value_message
   else:
     raise TypeError('Unsupported Python argument type {}.'.format(type(arg)))
 
@@ -77,6 +87,8 @@ def from_value_proto(result_pb):
     return result_pb.int_32
   elif which_result == 'float_32':
     return result_pb.float_32
+  elif which_result == 'custom_value':
+    return result_pb.custom_value
   elif which_result == 'struct':
     return [from_value_proto(x) for x in result_pb.struct.element]
   else:

genc/python/base/to_from_value_proto_test.py

@@ -14,6 +14,7 @@
 """Tests for to_from_value_proto.py."""
 
 from absl.testing import absltest
+from google3.google.protobuf import timestamp_pb2
 from genc.proto.v0 import computation_pb2 as pb
 from genc.python.base import to_from_value_proto
 
@@ -33,6 +34,16 @@ def test_something(self):
     self._roundtrip_test(0.5)
     self._roundtrip_test(["foo", 99])
 
+    # Test custom_value
+    timestamp_message = timestamp_pb2.Timestamp()
+    timestamp_message.FromJsonString("1970-01-01T00:00:00Z")
+    val_pb = to_from_value_proto.to_value_proto(timestamp_message)
+    self.assertIsInstance(val_pb, pb.Value)
+    any_val = to_from_value_proto.from_value_proto(val_pb)
+    reconstructed_timestamp_message = timestamp_pb2.Timestamp()
+    any_val.Unpack(reconstructed_timestamp_message)
+    self.assertEqual(timestamp_message, reconstructed_timestamp_message)
+
   def test_dict_to_value_proto(self):
     arg = {
         "key_1": "value_1",