Port gax proto helper methods (#4249)
diff --git a/google/api_core/protobuf_helpers.py b/google/api_core/protobuf_helpers.py
index 35eb574..6031ff0 100644
--- a/google/api_core/protobuf_helpers.py
+++ b/google/api_core/protobuf_helpers.py
@@ -14,6 +14,11 @@
"""Helpers for :mod:`protobuf`."""
+import collections
+import inspect
+
+from google.protobuf.message import Message
+
def from_any_pb(pb_type, any_pb):
"""Converts an ``Any`` protobuf to the specified message type.
@@ -36,3 +41,38 @@
any_pb.__class__.__name__, pb_type.__name__))
return msg
+
+
+def check_oneof(**kwargs):
+ """Raise ValueError if more than one keyword argument is not none.
+ Args:
+ kwargs (dict): The keyword arguments sent to the function.
+ Raises:
+ ValueError: If more than one entry in kwargs is not none.
+ """
+ # Sanity check: If no keyword arguments were sent, this is fine.
+ if not kwargs:
+ return
+
+ not_nones = [val for val in kwargs.values() if val is not None]
+ if len(not_nones) > 1:
+ raise ValueError('Only one of {fields} should be set.'.format(
+ fields=', '.join(sorted(kwargs.keys())),
+ ))
+
+
+def get_messages(module):
+ """Return a dictionary of message names and objects.
+ Args:
+ module (module): A Python module; dir() will be run against this
+ module to find Message subclasses.
+ Returns:
+ dict[str, Message]: A dictionary with the Message class names as
+ keys, and the Message subclasses themselves as values.
+ """
+ answer = collections.OrderedDict()
+ for name in dir(module):
+ candidate = getattr(module, name)
+ if inspect.isclass(candidate) and issubclass(candidate, Message):
+ answer[name] = candidate
+ return answer
diff --git a/tests/unit/test_protobuf_helpers.py b/tests/unit/test_protobuf_helpers.py
index 6233536..b9aca76 100644
--- a/tests/unit/test_protobuf_helpers.py
+++ b/tests/unit/test_protobuf_helpers.py
@@ -16,6 +16,7 @@
from google.api_core import protobuf_helpers
from google.protobuf import any_pb2
+from google.protobuf.message import Message
from google.type import date_pb2
from google.type import timeofday_pb2
@@ -35,3 +36,32 @@
with pytest.raises(TypeError):
protobuf_helpers.from_any_pb(timeofday_pb2.TimeOfDay, in_message)
+
+
+def test_check_protobuf_helpers_ok():
+ assert protobuf_helpers.check_oneof() is None
+ assert protobuf_helpers.check_oneof(foo='bar') is None
+ assert protobuf_helpers.check_oneof(foo='bar', baz=None) is None
+ assert protobuf_helpers.check_oneof(foo=None, baz='bacon') is None
+ assert (protobuf_helpers.check_oneof(foo='bar', spam=None, eggs=None)
+ is None)
+
+
+def test_check_protobuf_helpers_failures():
+ with pytest.raises(ValueError):
+ protobuf_helpers.check_oneof(foo='bar', spam='eggs')
+ with pytest.raises(ValueError):
+ protobuf_helpers.check_oneof(foo='bar', baz='bacon', spam='eggs')
+ with pytest.raises(ValueError):
+ protobuf_helpers.check_oneof(foo='bar', spam=0, eggs=None)
+
+
+def test_get_messages():
+ answer = protobuf_helpers.get_messages(date_pb2)
+
+ # Ensure that Date was exported properly.
+ assert answer['Date'] is date_pb2.Date
+
+ # Ensure that no non-Message objects were exported.
+ for value in answer.values():
+ assert issubclass(value, Message)