|
19 | 19 |
|
20 | 20 | import codecs
|
21 | 21 |
|
| 22 | +from collections.abc import Mapping as MappingABC |
| 23 | + |
22 | 24 |
|
23 | 25 | cdef uint64_t RECORD_ENCODER_CHECKED = 1 << 0
|
24 | 26 | cdef uint64_t RECORD_ENCODER_INVALID = 1 << 1
|
@@ -225,6 +227,86 @@ cdef class BaseNamedRecordCodec(BaseRecordCodec):
|
225 | 227 | (<BaseCodec>codec).dump(level + 1).strip()))
|
226 | 228 | return '\n'.join(buf)
|
227 | 229 |
|
| 230 | + cdef encode(self, WriteBuffer buf, object obj): |
| 231 | + cdef: |
| 232 | + WriteBuffer elem_data |
| 233 | + Py_ssize_t objlen |
| 234 | + Py_ssize_t i |
| 235 | + BaseCodec sub_codec |
| 236 | + Py_ssize_t is_dict |
| 237 | + Py_ssize_t is_namedtuple |
| 238 | + |
| 239 | + self._check_encoder() |
| 240 | + |
| 241 | + # We check in this order (dict, _is_array_iterable, |
| 242 | + # MappingABC) so that in the common case of dict or tuple, we |
| 243 | + # never do an ABC check. |
| 244 | + if cpython.PyDict_Check(obj): |
| 245 | + is_dict = True |
| 246 | + elif _is_array_iterable(obj): |
| 247 | + is_dict = False |
| 248 | + elif isinstance(obj, MappingABC): |
| 249 | + is_dict = True |
| 250 | + else: |
| 251 | + raise TypeError( |
| 252 | + 'a sized iterable container or mapping ' |
| 253 | + 'expected (got type {!r})'.format( |
| 254 | + type(obj).__name__)) |
| 255 | + is_namedtuple = not is_dict and hasattr(obj, '_fields') |
| 256 | + |
| 257 | + objlen = len(obj) |
| 258 | + if objlen == 0: |
| 259 | + buf.write_bytes(EMPTY_RECORD_DATA) |
| 260 | + return |
| 261 | + |
| 262 | + if objlen > _MAXINT32: |
| 263 | + raise ValueError('too many elements for a tuple') |
| 264 | + |
| 265 | + if objlen != len(self.fields_codecs): |
| 266 | + raise ValueError( |
| 267 | + f'expected {len(self.fields_codecs)} elements in the tuple, ' |
| 268 | + f'got {objlen}') |
| 269 | + |
| 270 | + elem_data = WriteBuffer.new() |
| 271 | + for i in range(objlen): |
| 272 | + if is_dict: |
| 273 | + name = datatypes.record_desc_pointer_name(self.descriptor, i) |
| 274 | + try: |
| 275 | + item = obj[name] |
| 276 | + except KeyError: |
| 277 | + raise ValueError( |
| 278 | + f"named tuple dict is missing '{name}' key", |
| 279 | + ) from None |
| 280 | + elif is_namedtuple: |
| 281 | + name = datatypes.record_desc_pointer_name(self.descriptor, i) |
| 282 | + try: |
| 283 | + item = getattr(obj, name) |
| 284 | + except AttributeError: |
| 285 | + raise ValueError( |
| 286 | + f"named tuple is missing '{name}' attribute", |
| 287 | + ) from None |
| 288 | + else: |
| 289 | + item = obj[i] |
| 290 | + |
| 291 | + elem_data.write_int32(0) # reserved bytes |
| 292 | + if item is None: |
| 293 | + elem_data.write_int32(-1) |
| 294 | + else: |
| 295 | + sub_codec = <BaseCodec>(self.fields_codecs[i]) |
| 296 | + try: |
| 297 | + sub_codec.encode(elem_data, item) |
| 298 | + except (TypeError, ValueError) as e: |
| 299 | + value_repr = repr(item) |
| 300 | + if len(value_repr) > 40: |
| 301 | + value_repr = value_repr[:40] + '...' |
| 302 | + raise errors.InvalidArgumentError( |
| 303 | + 'invalid input for query argument' |
| 304 | + ' ${n}: {v} ({msg})'.format( |
| 305 | + n=i, v=value_repr, msg=e)) from e |
| 306 | + |
| 307 | + buf.write_int32(4 + elem_data.len()) # buffer length |
| 308 | + buf.write_int32(<int32_t><uint32_t>objlen) |
| 309 | + buf.write_buffer(elem_data) |
228 | 310 |
|
229 | 311 | @cython.final
|
230 | 312 | cdef class EdegDBCodecContext(pgproto.CodecContext):
|
|
0 commit comments