Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/proto/marshal/marshal.py: 64%

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

110 statements  

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 

16 

17from google.protobuf import duration_pb2 

18from google.protobuf import timestamp_pb2 

19from google.protobuf import field_mask_pb2 

20from google.protobuf import struct_pb2 

21from google.protobuf import wrappers_pb2 

22 

23from proto.marshal import compat 

24from proto.marshal.collections import MapComposite 

25from proto.marshal.collections import Repeated 

26from proto.marshal.collections import RepeatedComposite 

27 

28from proto.marshal.rules import bytes as pb_bytes 

29from proto.marshal.rules import stringy_numbers 

30from proto.marshal.rules import dates 

31from proto.marshal.rules import struct 

32from proto.marshal.rules import wrappers 

33from proto.marshal.rules import field_mask 

34from proto.primitives import ProtoType 

35 

36 

37class Rule(abc.ABC): 

38 """Abstract class definition for marshal rules.""" 

39 

40 @classmethod 

41 def __subclasshook__(cls, C): 

42 if hasattr(C, "to_python") and hasattr(C, "to_proto"): 

43 return True 

44 return NotImplemented 

45 

46 

47class BaseMarshal: 

48 """The base class to translate between protobuf and Python classes. 

49 

50 Protocol buffers defines many common types (e.g. Timestamp, Duration) 

51 which also exist in the Python standard library. The marshal essentially 

52 translates between these: it keeps a registry of common protocol buffers 

53 and their Python representations, and translates back and forth. 

54 

55 The protocol buffer class is always the "key" in this relationship; when 

56 presenting a message, the declared field types are used to determine 

57 whether a value should be transformed into another class. Similarly, 

58 when accepting a Python value (when setting a field, for example), 

59 the declared field type is still used. This means that, if appropriate, 

60 multiple protocol buffer types may use the same Python type. 

61 

62 The primary implementation of this is :class:`Marshal`, which should 

63 usually be used instead of this class directly. 

64 """ 

65 

66 def __init__(self): 

67 self._rules = {} 

68 self._noop = NoopRule() 

69 self.reset() 

70 

71 def register(self, proto_type: type, rule: Rule = None): 

72 """Register a rule against the given ``proto_type``. 

73 

74 This function expects a ``proto_type`` (the descriptor class) and 

75 a ``rule``; an object with a ``to_python`` and ``to_proto`` method. 

76 Each method should return the appropriate Python or protocol buffer 

77 type, and be idempotent (e.g. accept either type as input). 

78 

79 This function can also be used as a decorator:: 

80 

81 @marshal.register(timestamp_pb2.Timestamp) 

82 class TimestampRule: 

83 ... 

84 

85 In this case, the class will be initialized for you with zero 

86 arguments. 

87 

88 Args: 

89 proto_type (type): A protocol buffer message type. 

90 rule: A marshal object 

91 """ 

92 # If a rule was provided, register it and be done. 

93 if rule: 

94 # Ensure the rule implements Rule. 

95 if not isinstance(rule, Rule): 

96 raise TypeError( 

97 "Marshal rule instances must implement " 

98 "`to_proto` and `to_python` methods." 

99 ) 

100 

101 # Register the rule. 

102 self._rules[proto_type] = rule 

103 return 

104 

105 # Create an inner function that will register an instance of the 

106 # marshal class to this object's registry, and return it. 

107 def register_rule_class(rule_class: type): 

108 # Ensure the rule class is a valid rule. 

109 if not issubclass(rule_class, Rule): 

110 raise TypeError( 

111 "Marshal rule subclasses must implement " 

112 "`to_proto` and `to_python` methods." 

113 ) 

114 

115 # Register the rule class. 

116 self._rules[proto_type] = rule_class() 

117 return rule_class 

118 

119 return register_rule_class 

120 

121 def reset(self): 

122 """Reset the registry to its initial state.""" 

123 self._rules.clear() 

124 

125 # Register date and time wrappers. 

126 self.register(timestamp_pb2.Timestamp, dates.TimestampRule()) 

127 self.register(duration_pb2.Duration, dates.DurationRule()) 

128 

129 # Register FieldMask wrappers. 

130 self.register(field_mask_pb2.FieldMask, field_mask.FieldMaskRule()) 

131 

132 # Register nullable primitive wrappers. 

133 self.register(wrappers_pb2.BoolValue, wrappers.BoolValueRule()) 

134 self.register(wrappers_pb2.BytesValue, wrappers.BytesValueRule()) 

135 self.register(wrappers_pb2.DoubleValue, wrappers.DoubleValueRule()) 

136 self.register(wrappers_pb2.FloatValue, wrappers.FloatValueRule()) 

137 self.register(wrappers_pb2.Int32Value, wrappers.Int32ValueRule()) 

138 self.register(wrappers_pb2.Int64Value, wrappers.Int64ValueRule()) 

139 self.register(wrappers_pb2.StringValue, wrappers.StringValueRule()) 

140 self.register(wrappers_pb2.UInt32Value, wrappers.UInt32ValueRule()) 

141 self.register(wrappers_pb2.UInt64Value, wrappers.UInt64ValueRule()) 

142 

143 # Register the google.protobuf.Struct wrappers. 

144 # 

145 # These are aware of the marshal that created them, because they 

146 # create RepeatedComposite and MapComposite instances directly and 

147 # need to pass the marshal to them. 

148 self.register(struct_pb2.Value, struct.ValueRule(marshal=self)) 

149 self.register(struct_pb2.ListValue, struct.ListValueRule(marshal=self)) 

150 self.register(struct_pb2.Struct, struct.StructRule(marshal=self)) 

151 

152 # Special case for bytes to allow base64 encode/decode 

153 self.register(ProtoType.BYTES, pb_bytes.BytesRule()) 

154 

155 # Special case for int64 from strings because of dict round trip. 

156 # See https://github.com/protocolbuffers/protobuf/issues/2679 

157 for rule_class in stringy_numbers.STRINGY_NUMBER_RULES: 

158 self.register(rule_class._proto_type, rule_class()) 

159 

160 def get_rule(self, proto_type): 

161 # Rules are needed to convert values between proto-plus and pb. 

162 # Retrieve the rule for the specified proto type. 

163 # The NoopRule will be used when a rule is not found. 

164 rule = self._rules.get(proto_type, self._noop) 

165 

166 # If we don't find a rule, also check under `_instances` 

167 # in case there is a rule in another package. 

168 # See https://github.com/googleapis/proto-plus-python/issues/349 

169 if rule == self._noop and hasattr(self, "_instances"): 

170 for _, instance in self._instances.items(): 

171 rule = instance._rules.get(proto_type, self._noop) 

172 if rule != self._noop: 

173 break 

174 return rule 

175 

176 def to_python(self, proto_type, value, *, absent: bool = None): 

177 # Internal protobuf has its own special type for lists of values. 

178 # Return a view around it that implements MutableSequence. 

179 value_type = type(value) # Minor performance boost over isinstance 

180 if value_type in compat.repeated_composite_types: 

181 return RepeatedComposite(value, marshal=self) 

182 if value_type in compat.repeated_scalar_types: 

183 if isinstance(proto_type, type): 

184 return RepeatedComposite(value, marshal=self, proto_type=proto_type) 

185 else: 

186 return Repeated(value, marshal=self) 

187 

188 # Same thing for maps of messages. 

189 # See https://github.com/protocolbuffers/protobuf/issues/16596 

190 # We need to look up the name of the type in compat.map_composite_type_names 

191 # as class `MessageMapContainer` is no longer exposed 

192 # This is done to avoid taking a breaking change in proto-plus. 

193 if ( 

194 value_type in compat.map_composite_types 

195 or value_type.__name__ in compat.map_composite_type_names 

196 ): 

197 return MapComposite(value, marshal=self) 

198 return self.get_rule(proto_type=proto_type).to_python(value, absent=absent) 

199 

200 def to_proto(self, proto_type, value, *, strict: bool = False): 

201 # The protos in google/protobuf/struct.proto are exceptional cases, 

202 # because they can and should represent themselves as lists and dicts. 

203 # These cases are handled in their rule classes. 

204 if proto_type not in ( 

205 struct_pb2.Value, 

206 struct_pb2.ListValue, 

207 struct_pb2.Struct, 

208 ): 

209 # For our repeated and map view objects, simply return the 

210 # underlying pb. 

211 if isinstance(value, (Repeated, MapComposite)): 

212 return value.pb 

213 

214 # Convert lists and tuples recursively. 

215 if isinstance(value, (list, tuple)): 

216 return type(value)(self.to_proto(proto_type, i) for i in value) 

217 

218 # Convert dictionaries recursively when the proto type is a map. 

219 # This is slightly more complicated than converting a list or tuple 

220 # because we have to step through the magic that protocol buffers does. 

221 # 

222 # Essentially, a type of map<string, Foo> will show up here as 

223 # a FoosEntry with a `key` field, `value` field, and a `map_entry` 

224 # annotation. We need to do the conversion based on the `value` 

225 # field's type. 

226 if isinstance(value, dict) and ( 

227 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()} 

232 

233 pb_value = self.get_rule(proto_type=proto_type).to_proto(value) 

234 

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 

246 

247 

248class Marshal(BaseMarshal): 

249 """The translator between protocol buffer and Python instances. 

250 

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 """ 

255 

256 _instances = {} 

257 

258 def __new__(cls, *, name: str): 

259 """Create a marshal instance. 

260 

261 Args: 

262 name (str): The name of the marshal. Instantiating multiple 

263 marshals with the same ``name`` argument will provide the 

264 same marshal each time. 

265 """ 

266 klass = cls._instances.get(name) 

267 if klass is None: 

268 klass = cls._instances[name] = super().__new__(cls) 

269 

270 return klass 

271 

272 def __init__(self, *, name: str): 

273 """Instantiate a marshal. 

274 

275 Args: 

276 name (str): The name of the marshal. Instantiating multiple 

277 marshals with the same ``name`` argument will provide the 

278 same marshal each time. 

279 """ 

280 self._name = name 

281 if not hasattr(self, "_rules"): 

282 super().__init__() 

283 

284 

285class NoopRule: 

286 """A catch-all rule that does nothing.""" 

287 

288 def to_python(self, pb_value, *, absent: bool = None): 

289 return pb_value 

290 

291 def to_proto(self, value): 

292 return value 

293 

294 

295__all__ = ("Marshal",)