Support pickling Record-s

Closes#451

Author: vmarkovtsev

asyncpg/protocol/protocol.pyx

@@ -1024,4 +1024,4 @@ def _create_record(object mapping, tuple elems):
 return rec
 
 
-Record=<object>record.ApgRecord_InitTypes()
+Record, RecordDescriptor =record.ApgRecord_InitTypes()

asyncpg/protocol/record/__init__.pxd

@@ -10,7 +10,7 @@ cimport cpython
 
 cdef extern from"record/recordobj.h":
 
-cpython.PyTypeObject *ApgRecord_InitTypes()exceptNULL
+tupleApgRecord_InitTypes()
 
 int ApgRecord_CheckExact(object)
 object ApgRecord_New(type, object, int)

asyncpg/protocol/record/recordobj.c

@@ -20,6 +20,11 @@ static PyObject * record_new_items_iter(PyObject *);
 staticApgRecordObject*free_list[ApgRecord_MAXSAVESIZE];
 staticintnumfree[ApgRecord_MAXSAVESIZE];
 
+staticPyObject*record_reconstruct_obj;
+staticPyObject*record_desc_reconstruct_obj;
+
+staticPyMethodDefrecord_desc_methods[];
+
 staticsize_tMAX_RECORD_SIZE= (
  ((size_t)PY_SSIZE_T_MAX-sizeof(ApgRecordObject) -sizeof(PyObject*))
  / sizeof(PyObject*)
@@ -575,14 +580,14 @@ record_repr(ApgRecordObject *v)
 
 
 staticPyObject*
-record_values(PyObject*o, PyObject*args)
+record_values(PyObject*o, PyObject*Py_UNUSED(unused))
 {
 returnrecord_iter(o);
 }
 
 
 staticPyObject*
-record_keys(PyObject*o, PyObject*args)
+record_keys(PyObject*o, PyObject*Py_UNUSED(unused))
 {
 if (!ApgRecord_Check(o)){
 PyErr_BadInternalCall();
@@ -594,7 +599,7 @@ record_keys(PyObject *o, PyObject *args)
 
 
 staticPyObject*
-record_items(PyObject*o, PyObject*args)
+record_items(PyObject*o, PyObject*Py_UNUSED(unused))
 {
 if (!ApgRecord_Check(o)){
 PyErr_BadInternalCall();
@@ -657,12 +662,71 @@ static PyMappingMethods record_as_mapping ={
 0/* mp_ass_subscript */
 };
 
+staticPyMethodDefrecord_methods[];
+
+staticPyObject*
+record_reduce(ApgRecordObject*o, PyObject*Py_UNUSED(unused))
+{
+PyObject*value=PyTuple_New(2);
+if (value==NULL){
+returnNULL;
+ }
+Py_ssize_tlen=Py_SIZE(o);
+PyObject*state=PyTuple_New(1+len);
+if (state==NULL){
+Py_DECREF(value);
+returnNULL;
+ }
+PyTuple_SET_ITEM(value, 0, record_reconstruct_obj);
+Py_INCREF(record_reconstruct_obj);
+PyTuple_SET_ITEM(value, 1, state);
+PyTuple_SET_ITEM(state, 0, (PyObject*)o->desc);
+Py_INCREF(o->desc);
+for (Py_ssize_ti=0; i<len; i++){
+PyObject*item=ApgRecord_GET_ITEM(o, i);
+PyTuple_SET_ITEM(state, i+1, item);
+Py_INCREF(item);
+ }
+returnvalue;
+}
+
+staticPyObject*
+record_reconstruct(PyObject*Py_UNUSED(unused), PyObject*args)
+{
+if (!PyTuple_CheckExact(args)){
+returnNULL;
+ }
+Py_ssize_tlen=PyTuple_GET_SIZE(args);
+if (len<2){
+returnNULL;
+ }
+len--;
+ApgRecordDescObject*desc= (ApgRecordDescObject*)PyTuple_GET_ITEM(args, 0);
+if (!ApgRecordDesc_CheckExact(desc)){
+returnNULL;
+ }
+if (PyObject_Length(desc->mapping) !=len){
+returnNULL;
+ }
+PyObject*record=ApgRecord_New(&ApgRecord_Type, (PyObject*)desc, len);
+if (record==NULL){
+returnNULL;
+ }
+for (Py_ssize_ti=0; i<len; i++){
+PyObject*item=PyTuple_GET_ITEM(args, i+1);
+ApgRecord_SET_ITEM(record, i, item);
+Py_INCREF(item);
+ }
+returnrecord;
+}
 
 staticPyMethodDefrecord_methods[] ={
 {"values", (PyCFunction)record_values, METH_NOARGS},
 {"keys", (PyCFunction)record_keys, METH_NOARGS},
 {"items", (PyCFunction)record_items, METH_NOARGS},
 {"get", (PyCFunction)record_get, METH_VARARGS},
+{"__reduce__", (PyCFunction)record_reduce, METH_NOARGS},
+{"__reconstruct__", (PyCFunction)record_reconstruct, METH_VARARGS | METH_STATIC},
 {NULL, NULL} /* sentinel */
 };
 
@@ -942,7 +1006,7 @@ record_new_items_iter(PyObject *seq)
 }
 
 
-PyTypeObject*
+PyObject*
 ApgRecord_InitTypes(void)
 {
 if (PyType_Ready(&ApgRecord_Type) <0){
@@ -961,7 +1025,22 @@ ApgRecord_InitTypes(void)
 returnNULL;
  }
 
-return&ApgRecord_Type;
+record_reconstruct_obj=PyCFunction_New(
+&record_methods[5], (PyObject*)&ApgRecord_Type
+ );
+record_desc_reconstruct_obj=PyCFunction_New(
+&record_desc_methods[1], (PyObject*)&ApgRecordDesc_Type
+ );
+
+PyObject*types=PyTuple_New(2);
+if (types==NULL){
+returnNULL;
+ }
+PyTuple_SET_ITEM(types, 0, (PyObject*)&ApgRecord_Type);
+Py_INCREF(&ApgRecord_Type);
+PyTuple_SET_ITEM(types, 1, (PyObject*)&ApgRecordDesc_Type);
+Py_INCREF(&ApgRecordDesc_Type);
+returntypes;
 }
 
 
@@ -987,15 +1066,54 @@ record_desc_traverse(ApgRecordDescObject *o, visitproc visit, void *arg)
 }
 
 
+staticPyObject*record_desc_reduce(ApgRecordDescObject*o, PyObject*Py_UNUSED(unused))
+{
+PyObject*value=PyTuple_New(2);
+if (value==NULL){
+returnNULL;
+ }
+PyObject*state=PyTuple_New(2);
+if (state==NULL){
+Py_DECREF(value);
+returnNULL;
+ }
+PyTuple_SET_ITEM(value, 0, record_desc_reconstruct_obj);
+Py_INCREF(record_desc_reconstruct_obj);
+PyTuple_SET_ITEM(value, 1, state);
+PyTuple_SET_ITEM(state, 0, o->mapping);
+Py_INCREF(o->mapping);
+PyTuple_SET_ITEM(state, 1, o->keys);
+Py_INCREF(o->keys);
+returnvalue;
+}
+
+
+staticPyObject*record_desc_reconstruct(PyObject*Py_UNUSED(unused), PyObject*args)
+{
+if (PyTuple_GET_SIZE(args) !=2){
+returnNULL;
+ }
+returnApgRecordDesc_New(PyTuple_GET_ITEM(args, 0), PyTuple_GET_ITEM(args, 1));
+}
+
+
+staticPyMethodDefrecord_desc_methods[] ={
+{"__reduce__", (PyCFunction)record_desc_reduce, METH_NOARGS},
+{"__reconstruct__", (PyCFunction)record_desc_reconstruct, METH_VARARGS | METH_STATIC},
+{NULL, NULL} /* sentinel */
+};
+
+
 PyTypeObjectApgRecordDesc_Type={
 PyVarObject_HEAD_INIT(NULL, 0)
- .tp_name="RecordDescriptor",
+ .tp_name="asyncpg.protocol.protocol.RecordDescriptor",
  .tp_basicsize=sizeof(ApgRecordDescObject),
  .tp_dealloc= (destructor)record_desc_dealloc,
  .tp_getattro=PyObject_GenericGetAttr,
  .tp_flags=Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
  .tp_traverse= (traverseproc)record_desc_traverse,
  .tp_iter=PyObject_SelfIter,
+ .tp_methods=record_desc_methods,
 };
 
 

asyncpg/protocol/record/recordobj.h

@@ -46,7 +46,7 @@ extern PyTypeObject ApgRecordDesc_Type;
 #defineApgRecord_GET_ITEM(op, i) \
  (((ApgRecordObject *)(op))->ob_item[i])
 
-PyTypeObject*ApgRecord_InitTypes(void);
+PyObject*ApgRecord_InitTypes(void);
 PyObject*ApgRecord_New(PyTypeObject*, PyObject*, Py_ssize_t);
 PyObject*ApgRecordDesc_New(PyObject*, PyObject*);
 

tests/test_record.py

@@ -575,3 +575,11 @@ class MyRecordBad:
 'record_class is expected to be a subclass of asyncpg.Record',
  ):
 awaitself.connect(record_class=MyRecordBad)
+
+deftest_record_pickle(self):
+r=pickle.loads(pickle.dumps(Record(R_AB, (42, 43))))
+self.assertEqual(len(r), 2)
+self.assertEqual(r[0], 42)
+self.assertEqual(r[1], 43)
+self.assertEqual(r['a'], 42)
+self.assertEqual(r['b'], 43)