Changeset View
Changeset View
Standalone View
Standalone View
swh/core/api/serializers.py
| Show All 36 Lines | DECODERS = { | ||||
| 'timedelta': lambda d: datetime.timedelta(**d), | 'timedelta': lambda d: datetime.timedelta(**d), | ||||
| 'uuid': UUID, | 'uuid': UUID, | ||||
| # Only for JSON: | # Only for JSON: | ||||
| 'bytes': base64.b85decode, | 'bytes': base64.b85decode, | ||||
| } | } | ||||
| def encode_data_client(data: Any) -> bytes: | def encode_data_client(data: Any, extra_encoders=None) -> bytes: | ||||
| try: | try: | ||||
| return msgpack_dumps(data) | return msgpack_dumps(data, extra_encoders=extra_encoders) | ||||
| except OverflowError as e: | except OverflowError as e: | ||||
| raise ValueError('Limits were reached. Please, check your input.\n' + | raise ValueError('Limits were reached. Please, check your input.\n' + | ||||
| str(e)) | str(e)) | ||||
| def decode_response(response: Response) -> Any: | def decode_response(response: Response, extra_decoders=None) -> Any: | ||||
| content_type = response.headers['content-type'] | content_type = response.headers['content-type'] | ||||
| if content_type.startswith('application/x-msgpack'): | if content_type.startswith('application/x-msgpack'): | ||||
| r = msgpack_loads(response.content) | r = msgpack_loads(response.content, | ||||
| extra_decoders=extra_decoders) | |||||
| elif content_type.startswith('application/json'): | elif content_type.startswith('application/json'): | ||||
| r = json.loads(response.text, cls=SWHJSONDecoder) | r = json.loads(response.text, cls=SWHJSONDecoder, | ||||
| extra_decoders=extra_decoders) | |||||
| elif content_type.startswith('text/'): | elif content_type.startswith('text/'): | ||||
| r = response.text | r = response.text | ||||
| else: | else: | ||||
| raise ValueError('Wrong content type `%s` for API response' | raise ValueError('Wrong content type `%s` for API response' | ||||
| % content_type) | % content_type) | ||||
| return r | return r | ||||
| Show All 17 Lines | class SWHJSONEncoder(json.JSONEncoder): | ||||
| Caveats: Limitations in the JSONEncoder extension mechanism | Caveats: Limitations in the JSONEncoder extension mechanism | ||||
| prevent us from "escaping" dictionaries that only contain the | prevent us from "escaping" dictionaries that only contain the | ||||
| swhtype and d keys, and therefore arbitrary data structures can't | swhtype and d keys, and therefore arbitrary data structures can't | ||||
| be round-tripped through SWHJSONEncoder and SWHJSONDecoder. | be round-tripped through SWHJSONEncoder and SWHJSONDecoder. | ||||
| """ | """ | ||||
| def __init__(self, extra_encoders=None, **kwargs): | |||||
| super().__init__(**kwargs) | |||||
| self.encoders = ENCODERS | |||||
| if extra_encoders: | |||||
| self.encoders += extra_encoders | |||||
| def default(self, o: Any | def default(self, o: Any | ||||
| ) -> Union[Dict[str, Union[Dict[str, int], str]], list]: | ) -> Union[Dict[str, Union[Dict[str, int], str]], list]: | ||||
| for (type_, type_name, encoder) in ENCODERS: | for (type_, type_name, encoder) in self.encoders: | ||||
| if isinstance(o, type_): | if isinstance(o, type_): | ||||
| return { | return { | ||||
| 'swhtype': type_name, | 'swhtype': type_name, | ||||
| 'd': encoder(o), | 'd': encoder(o), | ||||
| } | } | ||||
| try: | try: | ||||
| return super().default(o) | return super().default(o) | ||||
| except TypeError as e: | except TypeError as e: | ||||
| Show All 20 Lines | class SWHJSONDecoder(json.JSONDecoder): | ||||
| - swhtype with value 'bytes' or 'datetime'; | - swhtype with value 'bytes' or 'datetime'; | ||||
| - d containing the encoded value. | - d containing the encoded value. | ||||
| To limit the impact our encoding, if the swhtype key doesn't | To limit the impact our encoding, if the swhtype key doesn't | ||||
| contain a known value, the dictionary is decoded as-is. | contain a known value, the dictionary is decoded as-is. | ||||
| """ | """ | ||||
| def __init__(self, extra_decoders=None, **kwargs): | |||||
| super().__init__(**kwargs) | |||||
| self.decoders = DECODERS | |||||
| if extra_decoders: | |||||
| self.decoders = {**self.decoders, **extra_decoders} | |||||
| def decode_data(self, o: Any) -> Any: | def decode_data(self, o: Any) -> Any: | ||||
| if isinstance(o, dict): | if isinstance(o, dict): | ||||
| if set(o.keys()) == {'d', 'swhtype'}: | if set(o.keys()) == {'d', 'swhtype'}: | ||||
| if o['swhtype'] == 'bytes': | if o['swhtype'] == 'bytes': | ||||
| return base64.b85decode(o['d']) | return base64.b85decode(o['d']) | ||||
| decoder = DECODERS.get(o['swhtype']) | decoder = self.decoders.get(o['swhtype']) | ||||
| if decoder: | if decoder: | ||||
| return decoder(self.decode_data(o['d'])) | return decoder(self.decode_data(o['d'])) | ||||
| return {key: self.decode_data(value) for key, value in o.items()} | return {key: self.decode_data(value) for key, value in o.items()} | ||||
| if isinstance(o, list): | if isinstance(o, list): | ||||
| return [self.decode_data(value) for value in o] | return [self.decode_data(value) for value in o] | ||||
| else: | else: | ||||
| return o | return o | ||||
| def raw_decode(self, s: str, idx: int = 0) -> Tuple[Any, int]: | def raw_decode(self, s: str, idx: int = 0) -> Tuple[Any, int]: | ||||
| data, index = super().raw_decode(s, idx) | data, index = super().raw_decode(s, idx) | ||||
| return self.decode_data(data), index | return self.decode_data(data), index | ||||
| def msgpack_dumps(data: Any) -> bytes: | def msgpack_dumps(data: Any, extra_encoders=None) -> bytes: | ||||
| """Write data as a msgpack stream""" | """Write data as a msgpack stream""" | ||||
| encoders = ENCODERS | |||||
| if extra_encoders: | |||||
| encoders += extra_encoders | |||||
| def encode_types(obj): | def encode_types(obj): | ||||
| if isinstance(obj, types.GeneratorType): | if isinstance(obj, types.GeneratorType): | ||||
| return list(obj) | return list(obj) | ||||
| for (type_, type_name, encoder) in ENCODERS: | for (type_, type_name, encoder) in encoders: | ||||
| if isinstance(obj, type_): | if isinstance(obj, type_): | ||||
| return { | return { | ||||
| b'swhtype': type_name, | b'swhtype': type_name, | ||||
| b'd': encoder(obj), | b'd': encoder(obj), | ||||
| } | } | ||||
| return obj | return obj | ||||
| return msgpack.packb(data, use_bin_type=True, default=encode_types) | return msgpack.packb(data, use_bin_type=True, default=encode_types) | ||||
| def msgpack_loads(data: bytes) -> Any: | def msgpack_loads(data: bytes, extra_decoders=None) -> Any: | ||||
| """Read data as a msgpack stream""" | """Read data as a msgpack stream""" | ||||
| decoders = DECODERS | |||||
| if extra_decoders: | |||||
| decoders = {**decoders, **extra_decoders} | |||||
| def decode_types(obj): | def decode_types(obj): | ||||
| if set(obj.keys()) == {b'd', b'swhtype'}: | if set(obj.keys()) == {b'd', b'swhtype'}: | ||||
| decoder = DECODERS.get(obj[b'swhtype']) | decoder = decoders.get(obj[b'swhtype']) | ||||
| if decoder: | if decoder: | ||||
| return decoder(obj[b'd']) | return decoder(obj[b'd']) | ||||
| return obj | return obj | ||||
| try: | try: | ||||
| return msgpack.unpackb(data, raw=False, | return msgpack.unpackb(data, raw=False, | ||||
| object_hook=decode_types) | object_hook=decode_types) | ||||
| except TypeError: # msgpack < 0.5.2 | except TypeError: # msgpack < 0.5.2 | ||||
| return msgpack.unpackb(data, encoding='utf-8', | return msgpack.unpackb(data, encoding='utf-8', | ||||
| object_hook=decode_types) | object_hook=decode_types) | ||||