Skip to content

Commit 8c02697

Browse files
paulrosspaulross
committed
Add cCustomPickle. This is failing as can't import cPickle.
1 parent 592f955 commit 8c02697

File tree

4 files changed

+287
-1
lines changed

4 files changed

+287
-1
lines changed

‎CMakeLists.txt

+1-1
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ add_executable(PythonExtensionPatterns
9090
src/scratch.c
9191
# Legacy code. Removed in the documentation for version 0.2.0.
9292
# src/cpy/cParseArgsHelper.cpp
93-
)
93+
src/cpy/Pickle/cCustomPickle.c)
9494

9595
link_directories(
9696
${PYTHON_LINK_LIBRARY}

‎setup.py

+6
Original file line numberDiff line numberDiff line change
@@ -109,5 +109,11 @@
109109
extra_compile_args=extra_compile_args_c,
110110
language='c',
111111
),
112+
Extension(f"{PACKAGE_NAME}.cPickle", sources=['src/cpy/Pickle/cCustomPickle.c', ],
113+
include_dirs=['/usr/local/include', ], # os.path.join(os.getcwd(), 'include'),],
114+
library_dirs=[os.getcwd(), ], # path to .a or .so file(s)
115+
extra_compile_args=extra_compile_args_c,
116+
language='c',
117+
),
112118
]
113119
)

‎src/cpy/Pickle/cCustomPickle.c

+260
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
//
2+
// custom_pickle.c
3+
// PythonExtensionPatterns
4+
//
5+
// Created by Paul Ross on 08/04/2021.
6+
// Copyright (c) 2021 Paul Ross. All rights reserved.
7+
//
8+
// This adds pickling to a standard custom object.
9+
10+
#define PY_SSIZE_T_CLEAN
11+
#include <Python.h>
12+
#include "structmember.h"
13+
14+
typedef struct {
15+
PyObject_HEAD
16+
PyObject *first; /* first name */
17+
PyObject *last; /* last name */
18+
int number;
19+
} CustomObject;
20+
21+
static void
22+
Custom_dealloc(CustomObject *self)
23+
{
24+
Py_XDECREF(self->first);
25+
Py_XDECREF(self->last);
26+
Py_TYPE(self)->tp_free((PyObject *) self);
27+
}
28+
29+
static PyObject *
30+
Custom_new(PyTypeObject *type, PyObject *Py_UNUSED(args), PyObject *Py_UNUSED(kwds))
31+
{
32+
CustomObject *self;
33+
self = (CustomObject *) type->tp_alloc(type, 0);
34+
if (self != NULL) {
35+
self->first = PyUnicode_FromString("");
36+
if (self->first == NULL) {
37+
Py_DECREF(self);
38+
return NULL;
39+
}
40+
self->last = PyUnicode_FromString("");
41+
if (self->last == NULL) {
42+
Py_DECREF(self);
43+
return NULL;
44+
}
45+
self->number = 0;
46+
}
47+
fprintf(stdout, "Custom_new() reference counts first %zu last %zu\n", Py_REFCNT(self->first), Py_REFCNT(self->last));
48+
return (PyObject *) self;
49+
}
50+
51+
static int
52+
Custom_init(CustomObject *self, PyObject *args, PyObject *kwds)
53+
{
54+
static char *kwlist[] = {"first", "last", "number", NULL};
55+
PyObject *first = NULL, *last = NULL, *tmp;
56+
57+
if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOi", kwlist,
58+
&first, &last,
59+
&self->number))
60+
return -1;
61+
62+
if (first) {
63+
tmp = self->first;
64+
Py_INCREF(first);
65+
self->first = first;
66+
Py_XDECREF(tmp);
67+
}
68+
if (last) {
69+
tmp = self->last;
70+
Py_INCREF(last);
71+
self->last = last;
72+
Py_XDECREF(tmp);
73+
}
74+
return 0;
75+
}
76+
77+
static PyMemberDef Custom_members[] = {
78+
{"first", T_OBJECT_EX, offsetof(CustomObject, first), 0,
79+
"first name"},
80+
{"last", T_OBJECT_EX, offsetof(CustomObject, last), 0,
81+
"last name"},
82+
{"number", T_INT, offsetof(CustomObject, number), 0,
83+
"custom number"},
84+
{NULL, 0, 0, 0, NULL} /* Sentinel */
85+
};
86+
87+
static PyObject *
88+
Custom_name(CustomObject *self, PyObject *Py_UNUSED(ignored))
89+
{
90+
if (self->first == NULL) {
91+
PyErr_SetString(PyExc_AttributeError, "first");
92+
return NULL;
93+
}
94+
if (self->last == NULL) {
95+
PyErr_SetString(PyExc_AttributeError, "last");
96+
return NULL;
97+
}
98+
return PyUnicode_FromFormat("%S %S", self->first, self->last);
99+
}
100+
101+
/* Pickle the object */
102+
static const char* PICKLE_VERSION_KEY = "_pickle_version";
103+
static int PICKLE_VERSION = 1;
104+
105+
static PyObject *
106+
Custom___getstate__(CustomObject *self, PyObject *Py_UNUSED(ignored)) {
107+
PyObject *ret = Py_BuildValue("{sOsOsisi}",
108+
"first", self->first,
109+
"last", self->last,
110+
"number", self->number,
111+
PICKLE_VERSION_KEY, PICKLE_VERSION);
112+
fprintf(stdout, "Custom___getstate__ returning type %s\n", Py_TYPE(ret)->tp_name);
113+
return ret;
114+
}
115+
116+
117+
static PyObject *
118+
Custom___setstate__(CustomObject *self, PyObject *state) {
119+
fprintf(stdout, "Custom___getstate__ getting type %s\n", Py_TYPE(state)->tp_name);
120+
PyObject *key, *value;
121+
Py_ssize_t pos = 0;
122+
123+
while (PyDict_Next(state, &pos, &key, &value)) {
124+
/* do something interesting with the values... */
125+
fprintf(stdout, "Types Key: %s Value: %s\n", Py_TYPE(key)->tp_name, Py_TYPE(value)->tp_name);
126+
fprintf(stdout, "Key ");
127+
PyObject_Print(key, stdout, Py_PRINT_RAW);
128+
fprintf(stdout, " = ");
129+
PyObject_Print(value, stdout, Py_PRINT_RAW);
130+
fprintf(stdout, "\n");
131+
}
132+
fprintf(stdout, "Initial reference counts first %zu last %zu\n", Py_REFCNT(self->first), Py_REFCNT(self->last));
133+
134+
// static char *kwlist[] = {"first", "last", "number", NULL};
135+
//
136+
// PyArg_ParseTupleAndKeywords(args, state, "OOi", kwlist, &self->first, &self->last, &self->number);
137+
138+
#if 0
139+
// PyObject *key = NULL;
140+
Py_DECREF(self->first);
141+
key = Py_BuildValue("s", "first");
142+
self->first = PyDict_GetItem(state, key);
143+
Py_DECREF(key);
144+
Py_INCREF(self->first);
145+
146+
Py_DECREF(self->last);
147+
key = Py_BuildValue("s", "last");
148+
self->last = PyDict_GetItem(state, key);
149+
Py_DECREF(key);
150+
Py_INCREF(self->last);
151+
152+
key = Py_BuildValue("s", "number");
153+
self->number = PyLong_AsLong(PyDict_GetItem(state, key));
154+
Py_DECREF(key);
155+
#endif
156+
if (!PyDict_CheckExact(state)) {
157+
PyErr_SetString(PyExc_ValueError, "Pickled object is not a dict.");
158+
return NULL;
159+
}
160+
/* Version check. */
161+
/* Borrowed reference but no need to increment as we create a C long from it. */
162+
PyObject *temp = PyDict_GetItemString(state, PICKLE_VERSION_KEY);
163+
if (temp == NULL) {
164+
/* PyDict_GetItemString does not set any error state so we have to. */
165+
PyErr_Format(PyExc_KeyError, "No \"%s\" in pickled dict.", PICKLE_VERSION_KEY);
166+
return NULL;
167+
}
168+
int pickle_version = (int) PyLong_AsLong(temp);
169+
if (pickle_version != PICKLE_VERSION) {
170+
PyErr_Format(PyExc_ValueError, "Pickle version mismatch. Got version %d but expected version %d.",
171+
pickle_version, PICKLE_VERSION);
172+
return NULL;
173+
}
174+
175+
Py_DECREF(self->first);
176+
self->first = PyDict_GetItemString(state, "first"); /* Borrowed reference. */
177+
if (self->first == NULL) {
178+
/* PyDict_GetItemString does not set any error state so we have to. */
179+
PyErr_SetString(PyExc_KeyError, "No \"first\" in pickled dict.");
180+
return NULL;
181+
}
182+
Py_INCREF(self->first);
183+
184+
Py_DECREF(self->last);
185+
self->last = PyDict_GetItemString(state, "last"); /* Borrowed reference. */
186+
if (self->last == NULL) {
187+
/* PyDict_GetItemString does not set any error state so we have to. */
188+
PyErr_SetString(PyExc_KeyError, "No \"last\" in pickled dict.");
189+
return NULL;
190+
}
191+
Py_INCREF(self->last);
192+
193+
/* Borrowed reference but no need to increment as we create a C long from it. */
194+
PyObject *number = PyDict_GetItemString(state, "number");
195+
if (number == NULL) {
196+
/* PyDict_GetItemString does not set any error state so we have to. */
197+
PyErr_SetString(PyExc_KeyError, "No \"number\" in pickled dict.");
198+
return NULL;
199+
}
200+
self->number = (int) PyLong_AsLong(number);
201+
202+
fprintf(stdout, "Final reference counts first %zu last %zu\n", Py_REFCNT(self->first),
203+
Py_REFCNT(self->last));
204+
Py_RETURN_NONE;
205+
}
206+
207+
static PyMethodDef Custom_methods[] = {
208+
{"name", (PyCFunction) Custom_name, METH_NOARGS,
209+
"Return the name, combining the first and last name"
210+
},
211+
{"__getstate__", (PyCFunction) Custom___getstate__, METH_NOARGS,
212+
"Return the state for pickling"
213+
},
214+
{"__setstate__", (PyCFunction) Custom___setstate__, METH_O,
215+
"Set the state from a pickle"
216+
},
217+
{NULL, NULL, 0, NULL} /* Sentinel */
218+
};
219+
220+
static PyTypeObject CustomType = {
221+
PyVarObject_HEAD_INIT(NULL, 0)
222+
.tp_name = "cPickle.Custom",
223+
.tp_doc = "Custom objects",
224+
.tp_basicsize = sizeof(CustomObject),
225+
.tp_itemsize = 0,
226+
.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
227+
.tp_new = Custom_new,
228+
.tp_init = (initproc) Custom_init,
229+
.tp_dealloc = (destructor) Custom_dealloc,
230+
.tp_members = Custom_members,
231+
.tp_methods = Custom_methods,
232+
};
233+
234+
static PyModuleDef cPicklemodule = {
235+
PyModuleDef_HEAD_INIT,
236+
.m_name = "cPickle",
237+
.m_doc = "Example module that creates a pickleable extension type.",
238+
.m_size = -1,
239+
};
240+
241+
PyMODINIT_FUNC
242+
PyInit_cPickle(void)
243+
{
244+
PyObject *m;
245+
if (PyType_Ready(&CustomType) < 0)
246+
return NULL;
247+
248+
m = PyModule_Create(&cPicklemodule);
249+
if (m == NULL)
250+
return NULL;
251+
252+
Py_INCREF(&CustomType);
253+
if (PyModule_AddObject(m, "Custom", (PyObject *) &CustomType) < 0) {
254+
Py_DECREF(&CustomType);
255+
Py_DECREF(m);
256+
return NULL;
257+
}
258+
259+
return m;
260+
}

‎tests/unit/test_c_custom_pickle.py

+20
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
import pickle
2+
import pickletools
3+
import sys
4+
5+
import pytest
6+
7+
from cPyExtPatt import cPickle
8+
9+
10+
def test_module_dir():
11+
assert dir(cPickle) == ['Custom', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__spec__']
12+
13+
14+
def test_pickle_getstate():
15+
custom = cPickle.Custom('FIRST', 'LAST', 11)
16+
pickled_value = pickle.dumps(custom)
17+
print()
18+
print(f'Pickled original is {pickled_value}')
19+
assert pickled_value == b''
20+
# result = pickle.loads(pickled_value)

0 commit comments

Comments
 (0)