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