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

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 to_python(self, proto_type, value, *, absent: bool = None): 

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

164 # Return a view around it that implements MutableSequence. 

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

166 if value_type in compat.repeated_composite_types: 

167 return RepeatedComposite(value, marshal=self) 

168 if value_type in compat.repeated_scalar_types: 

169 if isinstance(proto_type, type): 

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

171 else: 

172 return Repeated(value, marshal=self) 

173 

174 # Same thing for maps of messages. 

175 if value_type in compat.map_composite_types: 

176 return MapComposite(value, marshal=self) 

177 

178 # Convert ordinary values. 

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

180 return rule.to_python(value, absent=absent) 

181 

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

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

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

185 # These cases are handled in their rule classes. 

186 if proto_type not in ( 

187 struct_pb2.Value, 

188 struct_pb2.ListValue, 

189 struct_pb2.Struct, 

190 ): 

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

192 # underlying pb. 

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

194 return value.pb 

195 

196 # Convert lists and tuples recursively. 

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

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

199 

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

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

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

203 # 

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

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

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

207 # field's type. 

208 if isinstance(value, dict) and ( 

209 proto_type.DESCRIPTOR.has_options 

210 and proto_type.DESCRIPTOR.GetOptions().map_entry 

211 ): 

212 recursive_type = type(proto_type().value) 

213 return {k: self.to_proto(recursive_type, v) for k, v in value.items()} 

214 

215 # Convert ordinary values. 

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

217 pb_value = rule.to_proto(value) 

218 

219 # Sanity check: If we are in strict mode, did we get the value we want? 

220 if strict and not isinstance(pb_value, proto_type): 

221 raise TypeError( 

222 "Parameter must be instance of the same class; " 

223 "expected {expected}, got {got}".format( 

224 expected=proto_type.__name__, 

225 got=pb_value.__class__.__name__, 

226 ), 

227 ) 

228 # Return the final value. 

229 return pb_value 

230 

231 

232class Marshal(BaseMarshal): 

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

234 

235 The bulk of the implementation is in :class:`BaseMarshal`. This class 

236 adds identity tracking: multiple instantiations of :class:`Marshal` with 

237 the same name will provide the same instance. 

238 """ 

239 

240 _instances = {} 

241 

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

243 """Create a marshal instance. 

244 

245 Args: 

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

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

248 same marshal each time. 

249 """ 

250 klass = cls._instances.get(name) 

251 if klass is None: 

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

253 

254 return klass 

255 

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

257 """Instantiate a marshal. 

258 

259 Args: 

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

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

262 same marshal each time. 

263 """ 

264 self._name = name 

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

266 super().__init__() 

267 

268 

269class NoopRule: 

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

271 

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

273 return pb_value 

274 

275 def to_proto(self, value): 

276 return value 

277 

278 

279__all__ = ("Marshal",)