Add support to unwrap Anys into wrapped pb2 objects. (#7430)
This commit adds support for unwrapping wrapped
pb2 objects from Anys (needed for LRO support with
wrapping).
diff --git a/google/api_core/protobuf_helpers.py b/google/api_core/protobuf_helpers.py
index 9762984..aec6b36 100644
--- a/google/api_core/protobuf_helpers.py
+++ b/google/api_core/protobuf_helpers.py
@@ -55,13 +55,22 @@
TypeError: if the message could not be converted.
"""
msg = pb_type()
- if not any_pb.Unpack(msg):
+
+ # Unwrap proto-plus wrapped messages.
+ if callable(getattr(pb_type, "pb", None)):
+ msg_pb = pb_type.pb(msg)
+ else:
+ msg_pb = msg
+
+ # Unpack the Any object and populate the protobuf message instance.
+ if not any_pb.Unpack(msg_pb):
raise TypeError(
"Could not convert {} to {}".format(
any_pb.__class__.__name__, pb_type.__name__
)
)
+ # Done; return the message.
return msg
diff --git a/tests/unit/test_protobuf_helpers.py b/tests/unit/test_protobuf_helpers.py
index ec761a0..db97238 100644
--- a/tests/unit/test_protobuf_helpers.py
+++ b/tests/unit/test_protobuf_helpers.py
@@ -38,6 +38,29 @@
assert in_message == out_message
+def test_from_any_pb_wrapped_success():
+ # Declare a message class conforming to wrapped messages.
+ class WrappedDate(object):
+ def __init__(self, **kwargs):
+ self._pb = date_pb2.Date(**kwargs)
+
+ def __eq__(self, other):
+ return self._pb == other
+
+ @classmethod
+ def pb(cls, msg):
+ return msg._pb
+
+ # Run the same test as `test_from_any_pb_success`, but using the
+ # wrapped class.
+ in_message = date_pb2.Date(year=1990)
+ in_message_any = any_pb2.Any()
+ in_message_any.Pack(in_message)
+ out_message = protobuf_helpers.from_any_pb(WrappedDate, in_message_any)
+
+ assert out_message == in_message
+
+
def test_from_any_pb_failure():
in_message = any_pb2.Any()
in_message.Pack(date_pb2.Date(year=1990))