Allow duplicate field names in records (fixes #28)

Author: 1st1

asyncpg/protocol/codecs/base.pyx

@@ -35,9 +35,14 @@ cdef class Codec:
         self.py_decoder = py_decoder
         self.element_codec = element_codec
         self.element_type_oids = element_type_oids
-        self.element_names = element_names
         self.element_codecs = element_codecs
 
+        if element_names is not None:
+            self.element_names = record.ApgRecordDesc_New(
+                element_names, tuple(element_names))
+        else:
+            self.element_names = None
+
         if type == CODEC_C:
             self.encoder = <codec_encode_func>&self.encode_scalar
             self.decoder = <codec_decode_func>&self.decode_scalar

asyncpg/protocol/prepared_stmt.pxd

@@ -25,7 +25,7 @@ cdef class PreparedStatementState:
         tuple        args_codecs
 
         int16_t      cols_num
-        object       cols_mapping
+        object       cols_desc
         bint         have_text_cols
         tuple        rows_codecs
 

asyncpg/protocol/prepared_stmt.pyx

@@ -16,7 +16,7 @@ cdef class PreparedStatementState:
         self.row_desc = self.parameters_desc = None
         self.args_codecs = self.rows_codecs = None
         self.args_num = self.cols_num = 0
-        self.cols_mapping = None
+        self.cols_desc = None
         self.closed = False
         self.refs = 0
         self.buffer = FastReadBuffer.new()
@@ -134,20 +134,24 @@ cdef class PreparedStatementState:
 
     cdef _ensure_rows_decoder(self):
         cdef:
+            list cols_names
             object cols_mapping
             tuple row
             int oid
             Codec codec
             list codecs
 
-        if self.cols_num == 0 or self.cols_mapping is not None:
+        if self.cols_num == 0 or self.cols_desc is not None:
             return
 
         cols_mapping = collections.OrderedDict()
+        cols_names = []
         codecs = []
         for i from 0 <= i < self.cols_num:
             row = self.row_desc[i]
-            cols_mapping[row[0].decode(self.settings._encoding)] = i
+            col_name = row[0].decode(self.settings._encoding)
+            cols_mapping[col_name] = i
+            cols_names.append(col_name)
             oid = row[3]
             codec = self.settings.get_data_codec(<uint32_t>oid)
             if codec is None or not codec.has_decoder():
@@ -157,7 +161,9 @@ cdef class PreparedStatementState:
 
             codecs.append(codec)
 
-        self.cols_mapping = cols_mapping
+        self.cols_desc = record.ApgRecordDesc_New(
+            cols_mapping, tuple(cols_names))
+
         self.rows_codecs = tuple(codecs)
 
     cdef _ensure_args_encoder(self):
@@ -215,7 +221,7 @@ cdef class PreparedStatementState:
         if rows_codecs is None or len(rows_codecs) < fnum:
             raise RuntimeError('invalid rows_codecs')
 
-        dec_row = record.ApgRecord_New(self.cols_mapping, fnum)
+        dec_row = record.ApgRecord_New(self.cols_desc, fnum)
         for i in range(fnum):
             flen = hton.unpack_int32(rbuf.read(4))
 

asyncpg/protocol/protocol.pyx

@@ -481,7 +481,8 @@ def _create_record(object mapping, tuple elems):
         object rec
         int32_t i
 
-    rec = record.ApgRecord_New(mapping, len(elems))
+    desc = record.ApgRecordDesc_New(mapping, tuple(mapping) if mapping else ())
+    rec = record.ApgRecord_New(desc, len(elems))
     for i in range(len(elems)):
         elem = elems[i]
         cpython.Py_INCREF(elem)

asyncpg/protocol/record/__init__.pxd

@@ -12,3 +12,5 @@ cdef extern from "record/recordobj.h":
 	int ApgRecord_CheckExact(object)
 	object ApgRecord_New(object, int)
 	void ApgRecord_SET_ITEM(object, int, object)
+
+	object ApgRecordDesc_New(object, object)

asyncpg/protocol/record/recordobj.c

@@ -17,12 +17,12 @@ static int numfree[ApgRecord_MAXSAVESIZE];
 
 
 PyObject *
-ApgRecord_New(PyObject *mapping, Py_ssize_t size)
+ApgRecord_New(PyObject *desc, Py_ssize_t size)
 {
     ApgRecordObject *o;
     Py_ssize_t i;
 
-    if (size < 1 || mapping == NULL) {
+    if (size < 1 || desc == NULL || !ApgRecordDesc_CheckExact(desc)) {
         PyErr_BadInternalCall();
         return NULL;
     }
@@ -48,8 +48,8 @@ ApgRecord_New(PyObject *mapping, Py_ssize_t size)
         o->ob_item[i] = NULL;
     }
 
-    Py_INCREF(mapping);
-    o->mapping = mapping;
+    Py_INCREF(desc);
+    o->desc = (ApgRecordDescObject*)desc;
     o->self_hash = -1;
     PyObject_GC_Track(o);
     return (PyObject *) o;
@@ -66,7 +66,7 @@ record_dealloc(ApgRecordObject *o)
 
     o->self_hash = -1;
 
-    Py_CLEAR(o->mapping);
+    Py_CLEAR(o->desc);
 
     Py_TRASHCAN_SAFE_BEGIN(o)
     if (len > 0) {
@@ -96,7 +96,7 @@ record_traverse(ApgRecordObject *o, visitproc visit, void *arg)
 {
     Py_ssize_t i;
 
-    Py_VISIT(o->mapping);
+    Py_VISIT(o->desc);
 
     for (i = Py_SIZE(o); --i >= 0;) {
         if (o->ob_item[i] != NULL) {
@@ -300,7 +300,7 @@ record_subscript(ApgRecordObject* o, PyObject* item)
     }
     else {
         PyObject *mapped;
-        mapped = PyObject_GetItem(o->mapping, item);
+        mapped = PyObject_GetItem(o->desc->mapping, item);
         if (mapped != NULL) {
             Py_ssize_t i;
             PyObject *result;
@@ -348,7 +348,7 @@ record_repr(ApgRecordObject *v)
     n = Py_SIZE(v);
     assert(n > 0);
 
-    keys_iter = PyObject_GetIter(v->mapping);
+    keys_iter = PyObject_GetIter(v->desc->keys);
     if (keys_iter == NULL) {
         return NULL;
     }
@@ -453,7 +453,7 @@ record_keys(PyObject *o, PyObject *args)
         return NULL;
     }
 
-    return PyObject_GetIter(((ApgRecordObject*)o)->mapping);
+    return PyObject_GetIter(((ApgRecordObject*)o)->desc->mapping);
 }
 
 
@@ -477,7 +477,7 @@ record_contains(ApgRecordObject *o, PyObject *arg)
         return -1;
     }
 
-    return PySequence_Contains(o->mapping, arg);
+    return PySequence_Contains(o->desc->mapping, arg);
 }
 
 
@@ -828,7 +828,7 @@ record_new_items_iter(PyObject *seq)
         return NULL;
     }
 
-    map_iter = PyObject_GetIter(((ApgRecordObject*)seq)->mapping);
+    map_iter = PyObject_GetIter(((ApgRecordObject*)seq)->desc->mapping);
     if (map_iter == NULL) {
         return NULL;
     }
@@ -853,6 +853,10 @@ int ApgRecord_InitTypes(void)
         return -1;
     }
 
+    if (PyType_Ready(&ApgRecordDesc_Type) < 0) {
+        return -1;
+    }
+
     if (PyType_Ready(&ApgRecordIter_Type) < 0) {
         return -1;
     }
@@ -863,3 +867,85 @@ int ApgRecord_InitTypes(void)
 
     return 0;
 }
+
+
+/* ----------------- */
+
+
+static void
+record_desc_dealloc(ApgRecordDescObject *o)
+{
+    PyObject_GC_UnTrack(o);
+    Py_CLEAR(o->mapping);
+    Py_CLEAR(o->keys);
+    PyObject_GC_Del(o);
+}
+
+
+static int
+record_desc_traverse(ApgRecordDescObject *o, visitproc visit, void *arg)
+{
+    Py_VISIT(o->mapping);
+    Py_VISIT(o->keys);
+    return 0;
+}
+
+
+PyTypeObject ApgRecordDesc_Type = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "RecordDescriptor",                         /* tp_name */
+    sizeof(ApgRecordDescObject),                /* tp_basicsize */
+    0,                                          /* tp_itemsize */
+    /* methods */
+    (destructor)record_desc_dealloc,            /* tp_dealloc */
+    0,                                          /* tp_print */
+    0,                                          /* tp_getattr */
+    0,                                          /* tp_setattr */
+    0,                                          /* tp_reserved */
+    0,                                          /* tp_repr */
+    0,                                          /* tp_as_number */
+    0,                                          /* tp_as_sequence */
+    0,                                          /* tp_as_mapping */
+    0,                                          /* tp_hash */
+    0,                                          /* tp_call */
+    0,                                          /* tp_str */
+    PyObject_GenericGetAttr,                    /* tp_getattro */
+    0,                                          /* tp_setattro */
+    0,                                          /* tp_as_buffer */
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
+    0,                                          /* tp_doc */
+    (traverseproc)record_desc_traverse,         /* tp_traverse */
+    0,                                          /* tp_clear */
+    0,                                          /* tp_richcompare */
+    0,                                          /* tp_weaklistoffset */
+    PyObject_SelfIter,                          /* tp_iter */
+    0,                                          /* tp_iternext */
+    0,                                          /* tp_methods */
+    0,
+};
+
+
+PyObject *
+ApgRecordDesc_New(PyObject *mapping, PyObject *keys)
+{
+    ApgRecordDescObject *o;
+
+    if (!mapping || !keys || !PyTuple_CheckExact(keys)) {
+        PyErr_BadInternalCall();
+        return NULL;
+    }
+
+    o = PyObject_GC_New(ApgRecordDescObject, &ApgRecordDesc_Type);
+    if (o == NULL) {
+        return NULL;
+    }
+
+    Py_INCREF(mapping);
+    o->mapping = mapping;
+
+    Py_INCREF(keys);
+    o->keys = keys;
+
+    PyObject_GC_Track(o);
+    return (PyObject *) o;
+}

asyncpg/protocol/record/recordobj.h

@@ -11,10 +11,17 @@
 #define ApgRecord_MAXFREELIST 2000
 
 
+typedef struct {
+    PyObject_HEAD
+    PyObject *mapping;
+    PyObject *keys;
+} ApgRecordDescObject;
+
+
 typedef struct {
     PyObject_VAR_HEAD
     Py_hash_t self_hash;
-    PyObject *mapping;
+    ApgRecordDescObject *desc;
     PyObject *ob_item[1];
 
     /* ob_item contains space for 'ob_size' elements.
@@ -23,17 +30,23 @@ typedef struct {
      */
 } ApgRecordObject;
 
+
 extern PyTypeObject ApgRecord_Type;
 extern PyTypeObject ApgRecordIter_Type;
 extern PyTypeObject ApgRecordItems_Type;
 
+extern PyTypeObject ApgRecordDesc_Type;
+
 #define ApgRecord_CheckExact(o) (Py_TYPE(o) == &ApgRecord_Type)
+#define ApgRecordDesc_CheckExact(o) (Py_TYPE(o) == &ApgRecordDesc_Type)
+
 #define ApgRecord_SET_ITEM(op, i, v) \
 			(((ApgRecordObject *)(op))->ob_item[i] = v)
 #define ApgRecord_GET_ITEM(op, i) \
 			(((ApgRecordObject *)(op))->ob_item[i])
 
 int ApgRecord_InitTypes(void);
 PyObject * ApgRecord_New(PyObject *, Py_ssize_t);
+PyObject * ApgRecordDesc_New(PyObject *, PyObject *);
 
 #endif

tests/test_record.py

@@ -281,7 +281,6 @@ def test_record_not_pickleable(self):
         with self.assertRaises(Exception):
             pickle.dumps(r)
 
-    @unittest.expectedFailure
     async def test_record_duplicate_colnames(self):
         """Test that Record handles duplicate column names."""
         r = await self.con.fetchrow('SELECT 1 as a, 2 as a')