OLD | NEW |
(Empty) | |
| 1 #!/usr/bin/env python |
| 2 # |
| 3 # Copyright 2015 Google Inc. |
| 4 # |
| 5 # Licensed under the Apache License, Version 2.0 (the "License"); |
| 6 # you may not use this file except in compliance with the License. |
| 7 # You may obtain a copy of the License at |
| 8 # |
| 9 # http://www.apache.org/licenses/LICENSE-2.0 |
| 10 # |
| 11 # Unless required by applicable law or agreed to in writing, software |
| 12 # distributed under the License is distributed on an "AS IS" BASIS, |
| 13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 # See the License for the specific language governing permissions and |
| 15 # limitations under the License. |
| 16 |
| 17 """Common code for converting proto to other formats, such as JSON.""" |
| 18 |
| 19 import base64 |
| 20 import collections |
| 21 import datetime |
| 22 import json |
| 23 import logging |
| 24 import os |
| 25 import sys |
| 26 |
| 27 import six |
| 28 |
| 29 from apitools.base.protorpclite import message_types |
| 30 from apitools.base.protorpclite import messages |
| 31 from apitools.base.protorpclite import protojson |
| 32 from apitools.base.py import exceptions |
| 33 |
| 34 __all__ = [ |
| 35 'CopyProtoMessage', |
| 36 'JsonToMessage', |
| 37 'MessageToJson', |
| 38 'DictToMessage', |
| 39 'MessageToDict', |
| 40 'PyValueToMessage', |
| 41 'MessageToPyValue', |
| 42 'MessageToRepr', |
| 43 'GetCustomJsonFieldMapping', |
| 44 'AddCustomJsonFieldMapping', |
| 45 'GetCustomJsonEnumMapping', |
| 46 'AddCustomJsonEnumMapping', |
| 47 ] |
| 48 |
| 49 |
| 50 _Codec = collections.namedtuple('_Codec', ['encoder', 'decoder']) |
| 51 CodecResult = collections.namedtuple('CodecResult', ['value', 'complete']) |
| 52 |
| 53 |
| 54 # TODO(craigcitro): Make these non-global. |
| 55 _UNRECOGNIZED_FIELD_MAPPINGS = {} |
| 56 _CUSTOM_MESSAGE_CODECS = {} |
| 57 _CUSTOM_FIELD_CODECS = {} |
| 58 _FIELD_TYPE_CODECS = {} |
| 59 |
| 60 |
| 61 def MapUnrecognizedFields(field_name): |
| 62 """Register field_name as a container for unrecognized fields.""" |
| 63 def Register(cls): |
| 64 _UNRECOGNIZED_FIELD_MAPPINGS[cls] = field_name |
| 65 return cls |
| 66 return Register |
| 67 |
| 68 |
| 69 def RegisterCustomMessageCodec(encoder, decoder): |
| 70 """Register a custom encoder/decoder for this message class.""" |
| 71 def Register(cls): |
| 72 _CUSTOM_MESSAGE_CODECS[cls] = _Codec(encoder=encoder, decoder=decoder) |
| 73 return cls |
| 74 return Register |
| 75 |
| 76 |
| 77 def RegisterCustomFieldCodec(encoder, decoder): |
| 78 """Register a custom encoder/decoder for this field.""" |
| 79 def Register(field): |
| 80 _CUSTOM_FIELD_CODECS[field] = _Codec(encoder=encoder, decoder=decoder) |
| 81 return field |
| 82 return Register |
| 83 |
| 84 |
| 85 def RegisterFieldTypeCodec(encoder, decoder): |
| 86 """Register a custom encoder/decoder for all fields of this type.""" |
| 87 def Register(field_type): |
| 88 _FIELD_TYPE_CODECS[field_type] = _Codec( |
| 89 encoder=encoder, decoder=decoder) |
| 90 return field_type |
| 91 return Register |
| 92 |
| 93 |
| 94 # TODO(craigcitro): Delete this function with the switch to proto2. |
| 95 def CopyProtoMessage(message): |
| 96 codec = protojson.ProtoJson() |
| 97 return codec.decode_message(type(message), codec.encode_message(message)) |
| 98 |
| 99 |
| 100 def MessageToJson(message, include_fields=None): |
| 101 """Convert the given message to JSON.""" |
| 102 result = _ProtoJsonApiTools.Get().encode_message(message) |
| 103 return _IncludeFields(result, message, include_fields) |
| 104 |
| 105 |
| 106 def JsonToMessage(message_type, message): |
| 107 """Convert the given JSON to a message of type message_type.""" |
| 108 return _ProtoJsonApiTools.Get().decode_message(message_type, message) |
| 109 |
| 110 |
| 111 # TODO(craigcitro): Do this directly, instead of via JSON. |
| 112 def DictToMessage(d, message_type): |
| 113 """Convert the given dictionary to a message of type message_type.""" |
| 114 return JsonToMessage(message_type, json.dumps(d)) |
| 115 |
| 116 |
| 117 def MessageToDict(message): |
| 118 """Convert the given message to a dictionary.""" |
| 119 return json.loads(MessageToJson(message)) |
| 120 |
| 121 |
| 122 def PyValueToMessage(message_type, value): |
| 123 """Convert the given python value to a message of type message_type.""" |
| 124 return JsonToMessage(message_type, json.dumps(value)) |
| 125 |
| 126 |
| 127 def MessageToPyValue(message): |
| 128 """Convert the given message to a python value.""" |
| 129 return json.loads(MessageToJson(message)) |
| 130 |
| 131 |
| 132 def MessageToRepr(msg, multiline=False, **kwargs): |
| 133 """Return a repr-style string for a protorpc message. |
| 134 |
| 135 protorpc.Message.__repr__ does not return anything that could be considered |
| 136 python code. Adding this function lets us print a protorpc message in such |
| 137 a way that it could be pasted into code later, and used to compare against |
| 138 other things. |
| 139 |
| 140 Args: |
| 141 msg: protorpc.Message, the message to be repr'd. |
| 142 multiline: bool, True if the returned string should have each field |
| 143 assignment on its own line. |
| 144 **kwargs: {str:str}, Additional flags for how to format the string. |
| 145 |
| 146 Known **kwargs: |
| 147 shortstrings: bool, True if all string values should be |
| 148 truncated at 100 characters, since when mocking the contents |
| 149 typically don't matter except for IDs, and IDs are usually |
| 150 less than 100 characters. |
| 151 no_modules: bool, True if the long module name should not be printed with |
| 152 each type. |
| 153 |
| 154 Returns: |
| 155 str, A string of valid python (assuming the right imports have been made) |
| 156 that recreates the message passed into this function. |
| 157 |
| 158 """ |
| 159 |
| 160 # TODO(jasmuth): craigcitro suggests a pretty-printer from apitools/gen. |
| 161 |
| 162 indent = kwargs.get('indent', 0) |
| 163 |
| 164 def IndentKwargs(kwargs): |
| 165 kwargs = dict(kwargs) |
| 166 kwargs['indent'] = kwargs.get('indent', 0) + 4 |
| 167 return kwargs |
| 168 |
| 169 if isinstance(msg, list): |
| 170 s = '[' |
| 171 for item in msg: |
| 172 if multiline: |
| 173 s += '\n' + ' ' * (indent + 4) |
| 174 s += MessageToRepr( |
| 175 item, multiline=multiline, **IndentKwargs(kwargs)) + ',' |
| 176 if multiline: |
| 177 s += '\n' + ' ' * indent |
| 178 s += ']' |
| 179 return s |
| 180 |
| 181 if isinstance(msg, messages.Message): |
| 182 s = type(msg).__name__ + '(' |
| 183 if not kwargs.get('no_modules'): |
| 184 s = msg.__module__ + '.' + s |
| 185 names = sorted([field.name for field in msg.all_fields()]) |
| 186 for name in names: |
| 187 field = msg.field_by_name(name) |
| 188 if multiline: |
| 189 s += '\n' + ' ' * (indent + 4) |
| 190 value = getattr(msg, field.name) |
| 191 s += field.name + '=' + MessageToRepr( |
| 192 value, multiline=multiline, **IndentKwargs(kwargs)) + ',' |
| 193 if multiline: |
| 194 s += '\n' + ' ' * indent |
| 195 s += ')' |
| 196 return s |
| 197 |
| 198 if isinstance(msg, six.string_types): |
| 199 if kwargs.get('shortstrings') and len(msg) > 100: |
| 200 msg = msg[:100] |
| 201 |
| 202 if isinstance(msg, datetime.datetime): |
| 203 |
| 204 class SpecialTZInfo(datetime.tzinfo): |
| 205 |
| 206 def __init__(self, offset): |
| 207 super(SpecialTZInfo, self).__init__() |
| 208 self.offset = offset |
| 209 |
| 210 def __repr__(self): |
| 211 s = 'TimeZoneOffset(' + repr(self.offset) + ')' |
| 212 if not kwargs.get('no_modules'): |
| 213 s = 'apitools.base.protorpclite.util.' + s |
| 214 return s |
| 215 |
| 216 msg = datetime.datetime( |
| 217 msg.year, msg.month, msg.day, msg.hour, msg.minute, msg.second, |
| 218 msg.microsecond, SpecialTZInfo(msg.tzinfo.utcoffset(0))) |
| 219 |
| 220 return repr(msg) |
| 221 |
| 222 |
| 223 def _GetField(message, field_path): |
| 224 for field in field_path: |
| 225 if field not in dir(message): |
| 226 raise KeyError('no field "%s"' % field) |
| 227 message = getattr(message, field) |
| 228 return message |
| 229 |
| 230 |
| 231 def _SetField(dictblob, field_path, value): |
| 232 for field in field_path[:-1]: |
| 233 dictblob[field] = {} |
| 234 dictblob = dictblob[field] |
| 235 dictblob[field_path[-1]] = value |
| 236 |
| 237 |
| 238 def _IncludeFields(encoded_message, message, include_fields): |
| 239 """Add the requested fields to the encoded message.""" |
| 240 if include_fields is None: |
| 241 return encoded_message |
| 242 result = json.loads(encoded_message) |
| 243 for field_name in include_fields: |
| 244 try: |
| 245 value = _GetField(message, field_name.split('.')) |
| 246 nullvalue = None |
| 247 if isinstance(value, list): |
| 248 nullvalue = [] |
| 249 except KeyError: |
| 250 raise exceptions.InvalidDataError( |
| 251 'No field named %s in message of type %s' % ( |
| 252 field_name, type(message))) |
| 253 _SetField(result, field_name.split('.'), nullvalue) |
| 254 return json.dumps(result) |
| 255 |
| 256 |
| 257 def _GetFieldCodecs(field, attr): |
| 258 result = [ |
| 259 getattr(_CUSTOM_FIELD_CODECS.get(field), attr, None), |
| 260 getattr(_FIELD_TYPE_CODECS.get(type(field)), attr, None), |
| 261 ] |
| 262 return [x for x in result if x is not None] |
| 263 |
| 264 |
| 265 class _ProtoJsonApiTools(protojson.ProtoJson): |
| 266 |
| 267 """JSON encoder used by apitools clients.""" |
| 268 _INSTANCE = None |
| 269 |
| 270 @classmethod |
| 271 def Get(cls): |
| 272 if cls._INSTANCE is None: |
| 273 cls._INSTANCE = cls() |
| 274 return cls._INSTANCE |
| 275 |
| 276 def decode_message(self, message_type, encoded_message): |
| 277 if message_type in _CUSTOM_MESSAGE_CODECS: |
| 278 return _CUSTOM_MESSAGE_CODECS[ |
| 279 message_type].decoder(encoded_message) |
| 280 # We turn off the default logging in protorpc. We may want to |
| 281 # remove this later. |
| 282 old_level = logging.getLogger().level |
| 283 logging.getLogger().setLevel(logging.ERROR) |
| 284 try: |
| 285 result = _DecodeCustomFieldNames(message_type, encoded_message) |
| 286 result = super(_ProtoJsonApiTools, self).decode_message( |
| 287 message_type, result) |
| 288 finally: |
| 289 logging.getLogger().setLevel(old_level) |
| 290 result = _ProcessUnknownEnums(result, encoded_message) |
| 291 result = _ProcessUnknownMessages(result, encoded_message) |
| 292 return _DecodeUnknownFields(result, encoded_message) |
| 293 |
| 294 def decode_field(self, field, value): |
| 295 """Decode the given JSON value. |
| 296 |
| 297 Args: |
| 298 field: a messages.Field for the field we're decoding. |
| 299 value: a python value we'd like to decode. |
| 300 |
| 301 Returns: |
| 302 A value suitable for assignment to field. |
| 303 """ |
| 304 for decoder in _GetFieldCodecs(field, 'decoder'): |
| 305 result = decoder(field, value) |
| 306 value = result.value |
| 307 if result.complete: |
| 308 return value |
| 309 if isinstance(field, messages.MessageField): |
| 310 field_value = self.decode_message( |
| 311 field.message_type, json.dumps(value)) |
| 312 elif isinstance(field, messages.EnumField): |
| 313 value = GetCustomJsonEnumMapping( |
| 314 field.type, json_name=value) or value |
| 315 try: |
| 316 field_value = super( |
| 317 _ProtoJsonApiTools, self).decode_field(field, value) |
| 318 except messages.DecodeError: |
| 319 if not isinstance(value, six.string_types): |
| 320 raise |
| 321 field_value = None |
| 322 else: |
| 323 field_value = super( |
| 324 _ProtoJsonApiTools, self).decode_field(field, value) |
| 325 return field_value |
| 326 |
| 327 def encode_message(self, message): |
| 328 if isinstance(message, messages.FieldList): |
| 329 return '[%s]' % (', '.join(self.encode_message(x) |
| 330 for x in message)) |
| 331 |
| 332 # pylint: disable=unidiomatic-typecheck |
| 333 if type(message) in _CUSTOM_MESSAGE_CODECS: |
| 334 return _CUSTOM_MESSAGE_CODECS[type(message)].encoder(message) |
| 335 |
| 336 message = _EncodeUnknownFields(message) |
| 337 result = super(_ProtoJsonApiTools, self).encode_message(message) |
| 338 result = _EncodeCustomFieldNames(message, result) |
| 339 return json.dumps(json.loads(result), sort_keys=True) |
| 340 |
| 341 def encode_field(self, field, value): |
| 342 """Encode the given value as JSON. |
| 343 |
| 344 Args: |
| 345 field: a messages.Field for the field we're encoding. |
| 346 value: a value for field. |
| 347 |
| 348 Returns: |
| 349 A python value suitable for json.dumps. |
| 350 """ |
| 351 for encoder in _GetFieldCodecs(field, 'encoder'): |
| 352 result = encoder(field, value) |
| 353 value = result.value |
| 354 if result.complete: |
| 355 return value |
| 356 if isinstance(field, messages.EnumField): |
| 357 if field.repeated: |
| 358 remapped_value = [GetCustomJsonEnumMapping( |
| 359 field.type, python_name=e.name) or e.name for e in value] |
| 360 else: |
| 361 remapped_value = GetCustomJsonEnumMapping( |
| 362 field.type, python_name=value.name) |
| 363 if remapped_value: |
| 364 return remapped_value |
| 365 if (isinstance(field, messages.MessageField) and |
| 366 not isinstance(field, message_types.DateTimeField)): |
| 367 value = json.loads(self.encode_message(value)) |
| 368 return super(_ProtoJsonApiTools, self).encode_field(field, value) |
| 369 |
| 370 |
| 371 # TODO(craigcitro): Fold this and _IncludeFields in as codecs. |
| 372 def _DecodeUnknownFields(message, encoded_message): |
| 373 """Rewrite unknown fields in message into message.destination.""" |
| 374 destination = _UNRECOGNIZED_FIELD_MAPPINGS.get(type(message)) |
| 375 if destination is None: |
| 376 return message |
| 377 pair_field = message.field_by_name(destination) |
| 378 if not isinstance(pair_field, messages.MessageField): |
| 379 raise exceptions.InvalidDataFromServerError( |
| 380 'Unrecognized fields must be mapped to a compound ' |
| 381 'message type.') |
| 382 pair_type = pair_field.message_type |
| 383 # TODO(craigcitro): Add more error checking around the pair |
| 384 # type being exactly what we suspect (field names, etc). |
| 385 if isinstance(pair_type.value, messages.MessageField): |
| 386 new_values = _DecodeUnknownMessages( |
| 387 message, json.loads(encoded_message), pair_type) |
| 388 else: |
| 389 new_values = _DecodeUnrecognizedFields(message, pair_type) |
| 390 setattr(message, destination, new_values) |
| 391 # We could probably get away with not setting this, but |
| 392 # why not clear it? |
| 393 setattr(message, '_Message__unrecognized_fields', {}) |
| 394 return message |
| 395 |
| 396 |
| 397 def _DecodeUnknownMessages(message, encoded_message, pair_type): |
| 398 """Process unknown fields in encoded_message of a message type.""" |
| 399 field_type = pair_type.value.type |
| 400 new_values = [] |
| 401 all_field_names = [x.name for x in message.all_fields()] |
| 402 for name, value_dict in six.iteritems(encoded_message): |
| 403 if name in all_field_names: |
| 404 continue |
| 405 value = PyValueToMessage(field_type, value_dict) |
| 406 if pair_type.value.repeated: |
| 407 value = _AsMessageList(value) |
| 408 new_pair = pair_type(key=name, value=value) |
| 409 new_values.append(new_pair) |
| 410 return new_values |
| 411 |
| 412 |
| 413 def _DecodeUnrecognizedFields(message, pair_type): |
| 414 """Process unrecognized fields in message.""" |
| 415 new_values = [] |
| 416 for unknown_field in message.all_unrecognized_fields(): |
| 417 # TODO(craigcitro): Consider validating the variant if |
| 418 # the assignment below doesn't take care of it. It may |
| 419 # also be necessary to check it in the case that the |
| 420 # type has multiple encodings. |
| 421 value, _ = message.get_unrecognized_field_info(unknown_field) |
| 422 value_type = pair_type.field_by_name('value') |
| 423 if isinstance(value_type, messages.MessageField): |
| 424 decoded_value = DictToMessage(value, pair_type.value.message_type) |
| 425 elif isinstance(value_type, messages.EnumField): |
| 426 decoded_value = pair_type.value.type(value) |
| 427 else: |
| 428 decoded_value = value |
| 429 new_pair = pair_type(key=str(unknown_field), value=decoded_value) |
| 430 new_values.append(new_pair) |
| 431 return new_values |
| 432 |
| 433 |
| 434 def _EncodeUnknownFields(message): |
| 435 """Remap unknown fields in message out of message.source.""" |
| 436 source = _UNRECOGNIZED_FIELD_MAPPINGS.get(type(message)) |
| 437 if source is None: |
| 438 return message |
| 439 result = CopyProtoMessage(message) |
| 440 pairs_field = message.field_by_name(source) |
| 441 if not isinstance(pairs_field, messages.MessageField): |
| 442 raise exceptions.InvalidUserInputError( |
| 443 'Invalid pairs field %s' % pairs_field) |
| 444 pairs_type = pairs_field.message_type |
| 445 value_variant = pairs_type.field_by_name('value').variant |
| 446 pairs = getattr(message, source) |
| 447 for pair in pairs: |
| 448 if value_variant == messages.Variant.MESSAGE: |
| 449 encoded_value = MessageToDict(pair.value) |
| 450 else: |
| 451 encoded_value = pair.value |
| 452 result.set_unrecognized_field(pair.key, encoded_value, value_variant) |
| 453 setattr(result, source, []) |
| 454 return result |
| 455 |
| 456 |
| 457 def _SafeEncodeBytes(field, value): |
| 458 """Encode the bytes in value as urlsafe base64.""" |
| 459 try: |
| 460 if field.repeated: |
| 461 result = [base64.urlsafe_b64encode(byte) for byte in value] |
| 462 else: |
| 463 result = base64.urlsafe_b64encode(value) |
| 464 complete = True |
| 465 except TypeError: |
| 466 result = value |
| 467 complete = False |
| 468 return CodecResult(value=result, complete=complete) |
| 469 |
| 470 |
| 471 def _SafeDecodeBytes(unused_field, value): |
| 472 """Decode the urlsafe base64 value into bytes.""" |
| 473 try: |
| 474 result = base64.urlsafe_b64decode(str(value)) |
| 475 complete = True |
| 476 except TypeError: |
| 477 result = value |
| 478 complete = False |
| 479 return CodecResult(value=result, complete=complete) |
| 480 |
| 481 |
| 482 def _ProcessUnknownEnums(message, encoded_message): |
| 483 """Add unknown enum values from encoded_message as unknown fields. |
| 484 |
| 485 ProtoRPC diverges from the usual protocol buffer behavior here and |
| 486 doesn't allow unknown fields. Throwing on unknown fields makes it |
| 487 impossible to let servers add new enum values and stay compatible |
| 488 with older clients, which isn't reasonable for us. We simply store |
| 489 unrecognized enum values as unknown fields, and all is well. |
| 490 |
| 491 Args: |
| 492 message: Proto message we've decoded thus far. |
| 493 encoded_message: JSON string we're decoding. |
| 494 |
| 495 Returns: |
| 496 message, with any unknown enums stored as unrecognized fields. |
| 497 """ |
| 498 if not encoded_message: |
| 499 return message |
| 500 decoded_message = json.loads(encoded_message) |
| 501 for field in message.all_fields(): |
| 502 if (isinstance(field, messages.EnumField) and |
| 503 field.name in decoded_message and |
| 504 message.get_assigned_value(field.name) is None): |
| 505 message.set_unrecognized_field( |
| 506 field.name, decoded_message[field.name], messages.Variant.ENUM) |
| 507 return message |
| 508 |
| 509 |
| 510 def _ProcessUnknownMessages(message, encoded_message): |
| 511 """Store any remaining unknown fields as strings. |
| 512 |
| 513 ProtoRPC currently ignores unknown values for which no type can be |
| 514 determined (and logs a "No variant found" message). For the purposes |
| 515 of reserializing, this is quite harmful (since it throws away |
| 516 information). Here we simply add those as unknown fields of type |
| 517 string (so that they can easily be reserialized). |
| 518 |
| 519 Args: |
| 520 message: Proto message we've decoded thus far. |
| 521 encoded_message: JSON string we're decoding. |
| 522 |
| 523 Returns: |
| 524 message, with any remaining unrecognized fields saved. |
| 525 """ |
| 526 if not encoded_message: |
| 527 return message |
| 528 decoded_message = json.loads(encoded_message) |
| 529 message_fields = [x.name for x in message.all_fields()] + list( |
| 530 message.all_unrecognized_fields()) |
| 531 missing_fields = [x for x in decoded_message.keys() |
| 532 if x not in message_fields] |
| 533 for field_name in missing_fields: |
| 534 message.set_unrecognized_field(field_name, decoded_message[field_name], |
| 535 messages.Variant.STRING) |
| 536 return message |
| 537 |
| 538 |
| 539 RegisterFieldTypeCodec(_SafeEncodeBytes, _SafeDecodeBytes)(messages.BytesField) |
| 540 |
| 541 |
| 542 # Note that these could share a dictionary, since they're keyed by |
| 543 # distinct types, but it's not really worth it. |
| 544 _JSON_ENUM_MAPPINGS = {} |
| 545 _JSON_FIELD_MAPPINGS = {} |
| 546 |
| 547 |
| 548 def _GetTypeKey(message_type, package): |
| 549 """Get the prefix for this message type in mapping dicts.""" |
| 550 key = message_type.definition_name() |
| 551 if package and key.startswith(package + '.'): |
| 552 module_name = message_type.__module__ |
| 553 # We normalize '__main__' to something unique, if possible. |
| 554 if module_name == '__main__': |
| 555 try: |
| 556 file_name = sys.modules[module_name].__file__ |
| 557 except (AttributeError, KeyError): |
| 558 pass |
| 559 else: |
| 560 base_name = os.path.basename(file_name) |
| 561 split_name = os.path.splitext(base_name) |
| 562 if len(split_name) == 1: |
| 563 module_name = unicode(base_name) |
| 564 else: |
| 565 module_name = u'.'.join(split_name[:-1]) |
| 566 key = module_name + '.' + key.partition('.')[2] |
| 567 return key |
| 568 |
| 569 |
| 570 def AddCustomJsonEnumMapping(enum_type, python_name, json_name, |
| 571 package=''): |
| 572 """Add a custom wire encoding for a given enum value. |
| 573 |
| 574 This is primarily used in generated code, to handle enum values |
| 575 which happen to be Python keywords. |
| 576 |
| 577 Args: |
| 578 enum_type: (messages.Enum) An enum type |
| 579 python_name: (basestring) Python name for this value. |
| 580 json_name: (basestring) JSON name to be used on the wire. |
| 581 package: (basestring, optional) Package prefix for this enum, if |
| 582 present. We strip this off the enum name in order to generate |
| 583 unique keys. |
| 584 """ |
| 585 if not issubclass(enum_type, messages.Enum): |
| 586 raise exceptions.TypecheckError( |
| 587 'Cannot set JSON enum mapping for non-enum "%s"' % enum_type) |
| 588 enum_name = _GetTypeKey(enum_type, package) |
| 589 if python_name not in enum_type.names(): |
| 590 raise exceptions.InvalidDataError( |
| 591 'Enum value %s not a value for type %s' % (python_name, enum_type)) |
| 592 field_mappings = _JSON_ENUM_MAPPINGS.setdefault(enum_name, {}) |
| 593 _CheckForExistingMappings('enum', enum_type, python_name, json_name) |
| 594 field_mappings[python_name] = json_name |
| 595 |
| 596 |
| 597 def AddCustomJsonFieldMapping(message_type, python_name, json_name, |
| 598 package=''): |
| 599 """Add a custom wire encoding for a given message field. |
| 600 |
| 601 This is primarily used in generated code, to handle enum values |
| 602 which happen to be Python keywords. |
| 603 |
| 604 Args: |
| 605 message_type: (messages.Message) A message type |
| 606 python_name: (basestring) Python name for this value. |
| 607 json_name: (basestring) JSON name to be used on the wire. |
| 608 package: (basestring, optional) Package prefix for this message, if |
| 609 present. We strip this off the message name in order to generate |
| 610 unique keys. |
| 611 """ |
| 612 if not issubclass(message_type, messages.Message): |
| 613 raise exceptions.TypecheckError( |
| 614 'Cannot set JSON field mapping for ' |
| 615 'non-message "%s"' % message_type) |
| 616 message_name = _GetTypeKey(message_type, package) |
| 617 try: |
| 618 _ = message_type.field_by_name(python_name) |
| 619 except KeyError: |
| 620 raise exceptions.InvalidDataError( |
| 621 'Field %s not recognized for type %s' % ( |
| 622 python_name, message_type)) |
| 623 field_mappings = _JSON_FIELD_MAPPINGS.setdefault(message_name, {}) |
| 624 _CheckForExistingMappings('field', message_type, python_name, json_name) |
| 625 field_mappings[python_name] = json_name |
| 626 |
| 627 |
| 628 def GetCustomJsonEnumMapping(enum_type, python_name=None, json_name=None): |
| 629 """Return the appropriate remapping for the given enum, or None.""" |
| 630 return _FetchRemapping(enum_type.definition_name(), 'enum', |
| 631 python_name=python_name, json_name=json_name, |
| 632 mappings=_JSON_ENUM_MAPPINGS) |
| 633 |
| 634 |
| 635 def GetCustomJsonFieldMapping(message_type, python_name=None, json_name=None): |
| 636 """Return the appropriate remapping for the given field, or None.""" |
| 637 return _FetchRemapping(message_type.definition_name(), 'field', |
| 638 python_name=python_name, json_name=json_name, |
| 639 mappings=_JSON_FIELD_MAPPINGS) |
| 640 |
| 641 |
| 642 def _FetchRemapping(type_name, mapping_type, python_name=None, json_name=None, |
| 643 mappings=None): |
| 644 """Common code for fetching a key or value from a remapping dict.""" |
| 645 if python_name and json_name: |
| 646 raise exceptions.InvalidDataError( |
| 647 'Cannot specify both python_name and json_name ' |
| 648 'for %s remapping' % mapping_type) |
| 649 if not (python_name or json_name): |
| 650 raise exceptions.InvalidDataError( |
| 651 'Must specify either python_name or json_name for %s remapping' % ( |
| 652 mapping_type,)) |
| 653 field_remappings = mappings.get(type_name, {}) |
| 654 if field_remappings: |
| 655 if python_name: |
| 656 return field_remappings.get(python_name) |
| 657 elif json_name: |
| 658 if json_name in list(field_remappings.values()): |
| 659 return [k for k in field_remappings |
| 660 if field_remappings[k] == json_name][0] |
| 661 return None |
| 662 |
| 663 |
| 664 def _CheckForExistingMappings(mapping_type, message_type, |
| 665 python_name, json_name): |
| 666 """Validate that no mappings exist for the given values.""" |
| 667 if mapping_type == 'field': |
| 668 getter = GetCustomJsonFieldMapping |
| 669 elif mapping_type == 'enum': |
| 670 getter = GetCustomJsonEnumMapping |
| 671 remapping = getter(message_type, python_name=python_name) |
| 672 if remapping is not None and remapping != json_name: |
| 673 raise exceptions.InvalidDataError( |
| 674 'Cannot add mapping for %s "%s", already mapped to "%s"' % ( |
| 675 mapping_type, python_name, remapping)) |
| 676 remapping = getter(message_type, json_name=json_name) |
| 677 if remapping is not None and remapping != python_name: |
| 678 raise exceptions.InvalidDataError( |
| 679 'Cannot add mapping for %s "%s", already mapped to "%s"' % ( |
| 680 mapping_type, json_name, remapping)) |
| 681 |
| 682 |
| 683 def _EncodeCustomFieldNames(message, encoded_value): |
| 684 message_name = type(message).definition_name() |
| 685 field_remappings = list(_JSON_FIELD_MAPPINGS.get(message_name, {}).items()) |
| 686 if field_remappings: |
| 687 decoded_value = json.loads(encoded_value) |
| 688 for python_name, json_name in field_remappings: |
| 689 if python_name in encoded_value: |
| 690 decoded_value[json_name] = decoded_value.pop(python_name) |
| 691 encoded_value = json.dumps(decoded_value) |
| 692 return encoded_value |
| 693 |
| 694 |
| 695 def _DecodeCustomFieldNames(message_type, encoded_message): |
| 696 message_name = message_type.definition_name() |
| 697 field_remappings = _JSON_FIELD_MAPPINGS.get(message_name, {}) |
| 698 if field_remappings: |
| 699 decoded_message = json.loads(encoded_message) |
| 700 for python_name, json_name in list(field_remappings.items()): |
| 701 if json_name in decoded_message: |
| 702 decoded_message[python_name] = decoded_message.pop(json_name) |
| 703 encoded_message = json.dumps(decoded_message) |
| 704 return encoded_message |
| 705 |
| 706 |
| 707 def _AsMessageList(msg): |
| 708 """Convert the provided list-as-JsonValue to a list.""" |
| 709 # This really needs to live in extra_types, but extra_types needs |
| 710 # to import this file to be able to register codecs. |
| 711 # TODO(craigcitro): Split out a codecs module and fix this ugly |
| 712 # import. |
| 713 from apitools.base.py import extra_types |
| 714 |
| 715 def _IsRepeatedJsonValue(msg): |
| 716 """Return True if msg is a repeated value as a JsonValue.""" |
| 717 if isinstance(msg, extra_types.JsonArray): |
| 718 return True |
| 719 if isinstance(msg, extra_types.JsonValue) and msg.array_value: |
| 720 return True |
| 721 return False |
| 722 |
| 723 if not _IsRepeatedJsonValue(msg): |
| 724 raise ValueError('invalid argument to _AsMessageList') |
| 725 if isinstance(msg, extra_types.JsonValue): |
| 726 msg = msg.array_value |
| 727 if isinstance(msg, extra_types.JsonArray): |
| 728 msg = msg.entries |
| 729 return msg |
OLD | NEW |