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 2018 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# https://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.
15import abc
16import threading
18from google.protobuf import (
19 duration_pb2,
20 field_mask_pb2,
21 struct_pb2,
22 timestamp_pb2,
23 wrappers_pb2,
24)
26from proto.marshal import compat
27from proto.marshal.collections import MapComposite, Repeated, RepeatedComposite
28from proto.marshal.rules import bytes as pb_bytes
29from proto.marshal.rules import dates, field_mask, stringy_numbers, struct, wrappers
30from proto.primitives import ProtoType
33class Rule(abc.ABC):
34 """Abstract class definition for marshal rules."""
36 @classmethod
37 def __subclasshook__(cls, C):
38 if hasattr(C, "to_python") and hasattr(C, "to_proto"):
39 return True
40 return NotImplemented
43class BaseMarshal:
44 """The base class to translate between protobuf and Python classes.
46 Protocol buffers defines many common types (e.g. Timestamp, Duration)
47 which also exist in the Python standard library. The marshal essentially
48 translates between these: it keeps a registry of common protocol buffers
49 and their Python representations, and translates back and forth.
51 The protocol buffer class is always the "key" in this relationship; when
52 presenting a message, the declared field types are used to determine
53 whether a value should be transformed into another class. Similarly,
54 when accepting a Python value (when setting a field, for example),
55 the declared field type is still used. This means that, if appropriate,
56 multiple protocol buffer types may use the same Python type.
58 The primary implementation of this is :class:`Marshal`, which should
59 usually be used instead of this class directly.
60 """
62 def __init__(self):
63 self._rules = {}
64 self._noop = NoopRule()
65 self.reset()
67 def register(self, proto_type: type, rule: Rule = None):
68 """Register a rule against the given ``proto_type``.
70 This function expects a ``proto_type`` (the descriptor class) and
71 a ``rule``; an object with a ``to_python`` and ``to_proto`` method.
72 Each method should return the appropriate Python or protocol buffer
73 type, and be idempotent (e.g. accept either type as input).
75 This function can also be used as a decorator::
77 @marshal.register(timestamp_pb2.Timestamp)
78 class TimestampRule:
79 ...
81 In this case, the class will be initialized for you with zero
82 arguments.
84 Args:
85 proto_type (type): A protocol buffer message type.
86 rule: A marshal object
87 """
88 # If a rule was provided, register it and be done.
89 if rule:
90 # Ensure the rule implements Rule.
91 if not isinstance(rule, Rule):
92 raise TypeError(
93 "Marshal rule instances must implement "
94 "`to_proto` and `to_python` methods."
95 )
97 # Register the rule.
98 self._rules[proto_type] = rule
99 return
101 # Create an inner function that will register an instance of the
102 # marshal class to this object's registry, and return it.
103 def register_rule_class(rule_class: type):
104 # Ensure the rule class is a valid rule.
105 if not issubclass(rule_class, Rule):
106 raise TypeError(
107 "Marshal rule subclasses must implement "
108 "`to_proto` and `to_python` methods."
109 )
111 # Register the rule class.
112 self._rules[proto_type] = rule_class()
113 return rule_class
115 return register_rule_class
117 def reset(self):
118 """Reset the registry to its initial state."""
119 self._rules.clear()
121 # Register date and time wrappers.
122 self.register(timestamp_pb2.Timestamp, dates.TimestampRule())
123 self.register(duration_pb2.Duration, dates.DurationRule())
125 # Register FieldMask wrappers.
126 self.register(field_mask_pb2.FieldMask, field_mask.FieldMaskRule())
128 # Register nullable primitive wrappers.
129 self.register(wrappers_pb2.BoolValue, wrappers.BoolValueRule())
130 self.register(wrappers_pb2.BytesValue, wrappers.BytesValueRule())
131 self.register(wrappers_pb2.DoubleValue, wrappers.DoubleValueRule())
132 self.register(wrappers_pb2.FloatValue, wrappers.FloatValueRule())
133 self.register(wrappers_pb2.Int32Value, wrappers.Int32ValueRule())
134 self.register(wrappers_pb2.Int64Value, wrappers.Int64ValueRule())
135 self.register(wrappers_pb2.StringValue, wrappers.StringValueRule())
136 self.register(wrappers_pb2.UInt32Value, wrappers.UInt32ValueRule())
137 self.register(wrappers_pb2.UInt64Value, wrappers.UInt64ValueRule())
139 # Register the google.protobuf.Struct wrappers.
140 #
141 # These are aware of the marshal that created them, because they
142 # create RepeatedComposite and MapComposite instances directly and
143 # need to pass the marshal to them.
144 self.register(struct_pb2.Value, struct.ValueRule(marshal=self))
145 self.register(struct_pb2.ListValue, struct.ListValueRule(marshal=self))
146 self.register(struct_pb2.Struct, struct.StructRule(marshal=self))
148 # Special case for bytes to allow base64 encode/decode
149 self.register(ProtoType.BYTES, pb_bytes.BytesRule())
151 # Special case for int64 from strings because of dict round trip.
152 # See https://github.com/protocolbuffers/protobuf/issues/2679
153 for rule_class in stringy_numbers.STRINGY_NUMBER_RULES:
154 self.register(rule_class._proto_type, rule_class())
156 def get_rule(self, proto_type):
157 # Rules are needed to convert values between proto-plus and pb.
158 # Retrieve the rule for the specified proto type.
159 # The NoopRule will be used when a rule is not found.
160 rule = self._rules.get(proto_type, self._noop)
162 # If we don't find a rule, also check under `_instances`
163 # in case there is a rule in another package.
164 # See https://github.com/googleapis/proto-plus-python/issues/349
165 if rule == self._noop and hasattr(self, "_instances"):
166 for _, instance in self._instances.items():
167 # Avoid race condition where instance is added to _instances
168 # but __init__ hasn't run yet.
169 rules = getattr(instance, "_rules", {})
170 rule = rules.get(proto_type, self._noop)
171 if rule != self._noop:
172 break
173 return rule
175 def to_python(self, proto_type, value, *, absent: bool = None):
176 # Internal protobuf has its own special type for lists of values.
177 # Return a view around it that implements MutableSequence.
178 value_type = type(value) # Minor performance boost over isinstance
179 if value_type in compat.repeated_composite_types:
180 return RepeatedComposite(value, marshal=self)
181 if value_type in compat.repeated_scalar_types:
182 if isinstance(proto_type, type):
183 return RepeatedComposite(value, marshal=self, proto_type=proto_type)
184 else:
185 return Repeated(value, marshal=self)
187 # Same thing for maps of messages.
188 # See https://github.com/protocolbuffers/protobuf/issues/16596
189 # We need to look up the name of the type in compat.map_composite_type_names
190 # as class `MessageMapContainer` is no longer exposed
191 # This is done to avoid taking a breaking change in proto-plus.
192 if (
193 value_type in compat.map_composite_types
194 or value_type.__name__ in compat.map_composite_type_names
195 ):
196 return MapComposite(value, marshal=self)
197 return self.get_rule(proto_type=proto_type).to_python(value, absent=absent)
199 def to_proto(self, proto_type, value, *, strict: bool = False):
200 # The protos in google/protobuf/struct.proto are exceptional cases,
201 # because they can and should represent themselves as lists and dicts.
202 # These cases are handled in their rule classes.
203 if proto_type not in (
204 struct_pb2.Value,
205 struct_pb2.ListValue,
206 struct_pb2.Struct,
207 ):
208 # For our repeated and map view objects, simply return the
209 # underlying pb.
210 if isinstance(value, (Repeated, MapComposite)):
211 return value.pb
213 # Convert lists and tuples recursively.
214 if isinstance(value, (list, tuple)):
215 return type(value)(self.to_proto(proto_type, i) for i in value)
217 # Convert dictionaries recursively when the proto type is a map.
218 # This is slightly more complicated than converting a list or tuple
219 # because we have to step through the magic that protocol buffers does.
220 #
221 # Essentially, a type of map<string, Foo> will show up here as
222 # a FoosEntry with a `key` field, `value` field, and a `map_entry`
223 # annotation. We need to do the conversion based on the `value`
224 # field's type.
225 if isinstance(value, dict) and (
226 hasattr(proto_type, "DESCRIPTOR")
227 and proto_type.DESCRIPTOR.has_options
228 and proto_type.DESCRIPTOR.GetOptions().map_entry
229 ):
230 recursive_type = type(proto_type().value)
231 return {k: self.to_proto(recursive_type, v) for k, v in value.items()}
233 pb_value = self.get_rule(proto_type=proto_type).to_proto(value)
235 # Sanity check: If we are in strict mode, did we get the value we want?
236 if strict and not isinstance(pb_value, proto_type):
237 raise TypeError(
238 "Parameter must be instance of the same class; "
239 "expected {expected}, got {got}".format(
240 expected=proto_type.__name__,
241 got=pb_value.__class__.__name__,
242 ),
243 )
244 # Return the final value.
245 return pb_value
248class Marshal(BaseMarshal):
249 """The translator between protocol buffer and Python instances.
251 The bulk of the implementation is in :class:`BaseMarshal`. This class
252 adds identity tracking: multiple instantiations of :class:`Marshal` with
253 the same name will provide the same instance.
254 """
256 _instances = {}
257 _instance_creation_lock = threading.Lock()
259 def __new__(cls, *, name: str):
260 """Create a marshal instance.
262 Args:
263 name (str): The name of the marshal. Instantiating multiple
264 marshals with the same ``name`` argument will provide the
265 same marshal each time.
266 """
267 klass = cls._instances.get(name)
268 if klass is None:
269 with cls._instance_creation_lock:
270 # Double check inside lock to confirm another thread hasn't
271 # created the instance while we were waiting for the lock.
272 klass = cls._instances.get(name)
273 if klass is None:
274 # Use Copy-on-Write to avoid 'RuntimeError: dictionary changed size during iteration'
275 # in BaseMarshal.get_rule. This allows other threads to iterate over the old
276 # dictionary safely while we replace it with a new one atomically.
277 new_instances = cls._instances.copy()
278 klass = super().__new__(cls)
279 new_instances[name] = klass
280 cls._instances = new_instances
282 return klass
284 def __init__(self, *, name: str):
285 """Instantiate a marshal.
287 Args:
288 name (str): The name of the marshal. Instantiating multiple
289 marshals with the same ``name`` argument will provide the
290 same marshal each time.
291 """
292 self._name = name
293 if not hasattr(self, "_rules"):
294 super().__init__()
297class NoopRule:
298 """A catch-all rule that does nothing."""
300 def to_python(self, pb_value, *, absent: bool = None):
301 return pb_value
303 def to_proto(self, value):
304 return value
307__all__ = ("Marshal",)