Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright 2017 Google LLC
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15"""Helpers for :mod:`protobuf`."""
17import collections
18import collections.abc
19import copy
20import inspect
22from google.protobuf import field_mask_pb2, message, wrappers_pb2
24_SENTINEL = object()
25_WRAPPER_TYPES = (
26 wrappers_pb2.BoolValue,
27 wrappers_pb2.BytesValue,
28 wrappers_pb2.DoubleValue,
29 wrappers_pb2.FloatValue,
30 wrappers_pb2.Int32Value,
31 wrappers_pb2.Int64Value,
32 wrappers_pb2.StringValue,
33 wrappers_pb2.UInt32Value,
34 wrappers_pb2.UInt64Value,
35)
38def from_any_pb(pb_type, any_pb):
39 """Converts an ``Any`` protobuf to the specified message type.
41 Args:
42 pb_type (type): the type of the message that any_pb stores an instance
43 of.
44 any_pb (google.protobuf.any_pb2.Any): the object to be converted.
46 Returns:
47 pb_type: An instance of the pb_type message.
49 Raises:
50 TypeError: if the message could not be converted.
51 """
52 msg = pb_type()
54 # Unwrap proto-plus wrapped messages.
55 if callable(getattr(pb_type, "pb", None)):
56 msg_pb = pb_type.pb(msg)
57 else:
58 msg_pb = msg
60 # Unpack the Any object and populate the protobuf message instance.
61 if not any_pb.Unpack(msg_pb):
62 raise TypeError(
63 f"Could not convert `{any_pb.TypeName()}` with underlying type `google.protobuf.any_pb2.Any` to `{msg_pb.DESCRIPTOR.full_name}`"
64 )
66 # Done; return the message.
67 return msg
70def check_oneof(**kwargs):
71 """Raise ValueError if more than one keyword argument is not ``None``.
73 Args:
74 kwargs (dict): The keyword arguments sent to the function.
76 Raises:
77 ValueError: If more than one entry in ``kwargs`` is not ``None``.
78 """
79 # Sanity check: If no keyword arguments were sent, this is fine.
80 if not kwargs:
81 return
83 not_nones = [val for val in kwargs.values() if val is not None]
84 if len(not_nones) > 1:
85 raise ValueError(
86 "Only one of {fields} should be set.".format(
87 fields=", ".join(sorted(kwargs.keys()))
88 )
89 )
92def get_messages(module):
93 """Discovers all protobuf Message classes in a given import module.
95 Args:
96 module (module): A Python module; :func:`dir` will be run against this
97 module to find Message subclasses.
99 Returns:
100 dict[str, google.protobuf.message.Message]: A dictionary with the
101 Message class names as keys, and the Message subclasses themselves
102 as values.
103 """
104 answer = collections.OrderedDict()
105 for name in dir(module):
106 candidate = getattr(module, name)
107 if inspect.isclass(candidate) and issubclass(candidate, message.Message):
108 answer[name] = candidate
109 return answer
112def _resolve_subkeys(key, separator="."):
113 """Resolve a potentially nested key.
115 If the key contains the ``separator`` (e.g. ``.``) then the key will be
116 split on the first instance of the subkey::
118 >>> _resolve_subkeys('a.b.c')
119 ('a', 'b.c')
120 >>> _resolve_subkeys('d|e|f', separator='|')
121 ('d', 'e|f')
123 If not, the subkey will be :data:`None`::
125 >>> _resolve_subkeys('foo')
126 ('foo', None)
128 Args:
129 key (str): A string that may or may not contain the separator.
130 separator (str): The namespace separator. Defaults to `.`.
132 Returns:
133 Tuple[str, str]: The key and subkey(s).
134 """
135 parts = key.split(separator, 1)
137 if len(parts) > 1:
138 return parts
139 else:
140 return parts[0], None
143def get(msg_or_dict, key, default=_SENTINEL):
144 """Retrieve a key's value from a protobuf Message or dictionary.
146 Args:
147 mdg_or_dict (Union[~google.protobuf.message.Message, Mapping]): the
148 object.
149 key (str): The key to retrieve from the object.
150 default (Any): If the key is not present on the object, and a default
151 is set, returns that default instead. A type-appropriate falsy
152 default is generally recommended, as protobuf messages almost
153 always have default values for unset values and it is not always
154 possible to tell the difference between a falsy value and an
155 unset one. If no default is set then :class:`KeyError` will be
156 raised if the key is not present in the object.
158 Returns:
159 Any: The return value from the underlying Message or dict.
161 Raises:
162 KeyError: If the key is not found. Note that, for unset values,
163 messages and dictionaries may not have consistent behavior.
164 TypeError: If ``msg_or_dict`` is not a Message or Mapping.
165 """
166 # We may need to get a nested key. Resolve this.
167 key, subkey = _resolve_subkeys(key)
169 # Attempt to get the value from the two types of objects we know about.
170 # If we get something else, complain.
171 if isinstance(msg_or_dict, message.Message):
172 answer = getattr(msg_or_dict, key, default)
173 elif isinstance(msg_or_dict, collections.abc.Mapping):
174 answer = msg_or_dict.get(key, default)
175 else:
176 raise TypeError(
177 "get() expected a dict or protobuf message, got {!r}.".format(
178 type(msg_or_dict)
179 )
180 )
182 # If the object we got back is our sentinel, raise KeyError; this is
183 # a "not found" case.
184 if answer is _SENTINEL:
185 raise KeyError(key)
187 # If a subkey exists, call this method recursively against the answer.
188 if subkey is not None and answer is not default:
189 return get(answer, subkey, default=default)
191 return answer
194def _set_field_on_message(msg, key, value):
195 """Set helper for protobuf Messages."""
196 # Attempt to set the value on the types of objects we know how to deal
197 # with.
198 if isinstance(value, (collections.abc.MutableSequence, tuple)):
199 # Clear the existing repeated protobuf message of any elements
200 # currently inside it.
201 while getattr(msg, key):
202 getattr(msg, key).pop()
204 # Write our new elements to the repeated field.
205 for item in value:
206 if isinstance(item, collections.abc.Mapping):
207 getattr(msg, key).add(**item)
208 else:
209 # protobuf's RepeatedCompositeContainer doesn't support
210 # append.
211 getattr(msg, key).extend([item])
212 elif isinstance(value, collections.abc.Mapping):
213 # Assign the dictionary values to the protobuf message.
214 for item_key, item_value in value.items():
215 set(getattr(msg, key), item_key, item_value)
216 elif isinstance(value, message.Message):
217 getattr(msg, key).CopyFrom(value)
218 else:
219 setattr(msg, key, value)
222def set(msg_or_dict, key, value):
223 """Set a key's value on a protobuf Message or dictionary.
225 Args:
226 msg_or_dict (Union[~google.protobuf.message.Message, Mapping]): the
227 object.
228 key (str): The key to set.
229 value (Any): The value to set.
231 Raises:
232 TypeError: If ``msg_or_dict`` is not a Message or dictionary.
233 """
234 # Sanity check: Is our target object valid?
235 if not isinstance(msg_or_dict, (collections.abc.MutableMapping, message.Message)):
236 raise TypeError(
237 "set() expected a dict or protobuf message, got {!r}.".format(
238 type(msg_or_dict)
239 )
240 )
242 # We may be setting a nested key. Resolve this.
243 basekey, subkey = _resolve_subkeys(key)
245 # If a subkey exists, then get that object and call this method
246 # recursively against it using the subkey.
247 if subkey is not None:
248 if isinstance(msg_or_dict, collections.abc.MutableMapping):
249 msg_or_dict.setdefault(basekey, {})
250 set(get(msg_or_dict, basekey), subkey, value)
251 return
253 if isinstance(msg_or_dict, collections.abc.MutableMapping):
254 msg_or_dict[key] = value
255 else:
256 _set_field_on_message(msg_or_dict, key, value)
259def setdefault(msg_or_dict, key, value):
260 """Set the key on a protobuf Message or dictionary to a given value if the
261 current value is falsy.
263 Because protobuf Messages do not distinguish between unset values and
264 falsy ones particularly well (by design), this method treats any falsy
265 value (e.g. 0, empty list) as a target to be overwritten, on both Messages
266 and dictionaries.
268 Args:
269 msg_or_dict (Union[~google.protobuf.message.Message, Mapping]): the
270 object.
271 key (str): The key on the object in question.
272 value (Any): The value to set.
274 Raises:
275 TypeError: If ``msg_or_dict`` is not a Message or dictionary.
276 """
277 if not get(msg_or_dict, key, default=None):
278 set(msg_or_dict, key, value)
281def field_mask(original, modified):
282 """Create a field mask by comparing two messages.
284 Args:
285 original (~google.protobuf.message.Message): the original message.
286 If set to None, this field will be interpreted as an empty
287 message.
288 modified (~google.protobuf.message.Message): the modified message.
289 If set to None, this field will be interpreted as an empty
290 message.
292 Returns:
293 google.protobuf.field_mask_pb2.FieldMask: field mask that contains
294 the list of field names that have different values between the two
295 messages. If the messages are equivalent, then the field mask is empty.
297 Raises:
298 ValueError: If the ``original`` or ``modified`` are not the same type.
299 """
300 if original is None and modified is None:
301 return field_mask_pb2.FieldMask()
303 if original is None and modified is not None:
304 original = copy.deepcopy(modified)
305 original.Clear()
307 if modified is None and original is not None:
308 modified = copy.deepcopy(original)
309 modified.Clear()
311 if not isinstance(original, type(modified)):
312 raise ValueError(
313 "expected that both original and modified should be of the "
314 'same type, received "{!r}" and "{!r}".'.format(
315 type(original), type(modified)
316 )
317 )
319 return field_mask_pb2.FieldMask(paths=_field_mask_helper(original, modified))
322def _field_mask_helper(original, modified, current=""):
323 answer = []
325 for name in original.DESCRIPTOR.fields_by_name:
326 field_path = _get_path(current, name)
328 original_val = getattr(original, name)
329 modified_val = getattr(modified, name)
331 if _is_message(original_val) or _is_message(modified_val):
332 if original_val != modified_val:
333 # Wrapper types do not need to include the .value part of the
334 # path.
335 if _is_wrapper(original_val) or _is_wrapper(modified_val):
336 answer.append(field_path)
337 elif not modified_val.ListFields():
338 answer.append(field_path)
339 else:
340 answer.extend(
341 _field_mask_helper(original_val, modified_val, field_path)
342 )
343 else:
344 if original_val != modified_val:
345 answer.append(field_path)
347 return answer
350def _get_path(current, name):
351 # gapic-generator-python appends underscores to field names
352 # that collide with python keywords.
353 # `_` is stripped away as it is not possible to
354 # natively define a field with a trailing underscore in protobuf.
355 # APIs will reject field masks if fields have trailing underscores.
356 # See https://github.com/googleapis/python-api-core/issues/227
357 name = name.rstrip("_")
358 if not current:
359 return name
360 return "%s.%s" % (current, name)
363def _is_message(value):
364 return isinstance(value, message.Message)
367def _is_wrapper(value):
368 return type(value) in _WRAPPER_TYPES