Port gax proto helper methods (#4249)

Author: landrito

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 @@ def from_any_pb(pb_type, any_pb):
                 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

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 @@ def test_from_any_pb_failure():
 
     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)