Implementing pickle interface for RPNExpr

python_rpnexpr.c

@@ -22,6 +22,11 @@ PyMethodDef RPNExpr_methods[] = {
 	{"eval", (PyCFunction)rpnexpr_eval, METH_FASTCALL, "Evaluate an expression"},
 	{"reset_stack", (PyCFunction)rpnexpr_reset_stack, METH_NOARGS,
 		"Reset stack memory storage (set all items to 0)"},
+	{"__getstate__", (PyCFunction)rpnexpr_getstate, METH_NOARGS,
+		"Pickling method. Return a bytes repr of tokenized expression \
+and the stack state."},
+	{"__setstate__", (PyCFunction)rpnexpr_setstate, METH_O,
+		"Unpickling method"},
 	{NULL} //Sentinel
 };
 
@@ -31,7 +36,7 @@ PyMemberDef RPNExpr_members[] = {
 
 PyTypeObject RPNExprType = {
 	PyVarObject_HEAD_INIT(NULL, 0)
-	"rpn.RPNExpr",                     /* tp_name */
+	"pyrpn.RPNExpr",                     /* tp_name */
 	sizeof(PyRPNExpr_t),                        /* tp_basicsize */
 	0,                                               /* tp_itemsize */
 	(destructor)rpnexpr_del, /* tp_dealloc */
@@ -149,13 +154,7 @@ int rpnexpr_init(PyObject *self, PyObject *args, PyObject *kwds)
 		return -1;
 	}
 
-	if(!stack_size)
-	{
-		expr_self->args = NULL;
-		return 0;
-	}
-
-	expr_self->args = malloc(sizeof(unsigned long) * args_count);
+	expr_self->args = malloc(sizeof(rpn_value_t) * args_count);
 	if(!expr_self->args)
 	{
 		snprintf(err_str, 256,
@@ -191,6 +190,226 @@ void rpnexpr_del(PyObject *self)
 	}
 }
 
+
+PyObject* rpnexpr_getstate(PyObject *self, PyObject *noargs)
+{
+	PyObject *res, *part;
+	PyRPNExpr_state_t resbuf;
+	PyRPNExpr_t *expr_self;
+	size_t total_sz;
+	char err_str[128];
+
+	expr_self = (PyRPNExpr_t*)self;
+
+	if(expr_self->rpn->state != RPN_READY)
+	{
+		snprintf(err_str, 128, 
+			"RPNExpr.__getstate__() instance in bad state : %s",
+			expr_self->rpn->state==RPN_ERROR?"error":"init");
+		PyErr_SetString(PyExc_ValueError, err_str);
+		return NULL;
+	}
+
+	total_sz = sizeof(PyRPNExpr_state_t);
+	total_sz += expr_self->rpn->toks.tokens_sz * sizeof(rpn_token_t);
+	total_sz += expr_self->rpn->stack_sz * sizeof(rpn_value_t);
+
+	resbuf.total_sz = total_sz;
+	resbuf.argc = expr_self->rpn->args_count;
+	resbuf.stack_sz = expr_self->rpn->stack_sz;
+	resbuf.token_sz = expr_self->rpn->toks.tokens_sz;
+
+	if(!(res = PyBytes_FromStringAndSize((char*)&resbuf, sizeof(resbuf))))
+	{
+		return NULL;
+	}
+	if(resbuf.stack_sz)
+	{
+		if(!(part=PyBytes_FromStringAndSize(
+				(char*)expr_self->rpn->stack,
+				sizeof(rpn_value_t) * resbuf.stack_sz)))
+		{
+			return NULL;
+		}
+		PyBytes_ConcatAndDel(&res, part);
+		if(!res)
+		{
+			return NULL;
+		}
+	}
+	if(resbuf.token_sz)
+	{
+		if(!(part=PyBytes_FromStringAndSize(
+				(char*)expr_self->rpn->toks.tokens,
+				sizeof(rpn_token_t) * resbuf.token_sz)))
+		{
+			return NULL;
+		}
+		PyBytes_ConcatAndDel(&res, part);
+		if(!res)
+		{
+			return NULL;
+		}
+	}
+	return res;
+}
+
+PyObject* rpnexpr_setstate(PyObject *self, PyObject *state_bytes)
+{
+	PyObject *tmp, *tmp2;
+	PyRPNExpr_state_t *state;
+	PyRPNExpr_t *expr_self;
+	rpn_value_t *stack;
+	rpn_tokenized_t toks;
+	const char *data;
+	size_t bsize, csize;
+	int err;
+	char err_str[256];
+	size_t i;
+
+	expr_self = (PyRPNExpr_t*)self;
+	
+	if(!PyBytes_Check(state_bytes)) /* Arg check */
+	{
+		tmp2 = NULL;
+		tmp = PyObject_Type(state_bytes);
+		if(tmp)
+		{
+			tmp2 = PyObject_Str(tmp);
+		}
+		if(tmp2)
+		{
+			snprintf(err_str, 128, 
+				"RPNExpr.__setstate__() expected a bytes as \
+argument but %s found",
+				PyUnicode_AsUTF8(tmp2));
+			PyErr_SetString(PyExc_ValueError, err_str);
+		}
+		else
+		{
+			PyErr_SetString(PyExc_RuntimeError,
+				"Failing to fetch arguments type will \
+generating exception message !");
+		}
+		return NULL;
+	}
+	if(expr_self->rpn || expr_self->args) /* checking instance state */
+	{
+		PyErr_SetString(PyExc_ValueError,
+			"RPNExpr.__getstate__() instance in bad state : \
+should not be initialized");
+		return NULL;
+	}
+	
+	/* Checking data size */
+	bsize = PyBytes_GET_SIZE(state_bytes);
+	data = PyBytes_AS_STRING(state_bytes);
+	state = (PyRPNExpr_state_t*)data;
+
+	if(bsize < sizeof(size_t))
+	{
+		PyErr_SetString(PyExc_ValueError, "Invalid argument");
+	}
+	if(bsize != state->total_sz)
+	{
+		snprintf(err_str, 128,
+			"RPNExpr.__setstate__() error : expected state to \
+contains %ld bytes but %ld found",
+			state->total_sz, bsize);
+		PyErr_SetString(PyExc_ValueError, err_str);
+	}
+
+	/* Checking data "integrity" */
+	csize = sizeof(PyRPNExpr_state_t) +
+		state->token_sz * sizeof(rpn_token_t) +
+		state->stack_sz * sizeof(rpn_value_t);
+	if(state->total_sz != csize)
+	{
+		snprintf(err_str, 128,
+			"RPNExpr.__setstate__() error : should have %ld as \
+total size but %ld found",
+			csize, state->total_sz);
+		PyErr_SetString(PyExc_ValueError, err_str);
+		return NULL;
+	}
+
+
+	/* Alloc and init rpn expression & args buffer*/
+	if(!(expr_self->rpn = malloc(sizeof(rpn_expr_t))))
+	{
+		err = errno;
+		snprintf(err_str, 128, 
+			"RPNExpr.__setstate__() failed to allocate memory for \
+rpn_expr_t : %s",
+			strerror(err));
+		PyErr_SetString(PyExc_MemoryError, err_str);
+		return NULL;
+	}
+	bzero(expr_self->rpn, sizeof(rpn_expr_t));
+	if(!(expr_self->args = malloc(sizeof(rpn_value_t)*state->argc)))
+	{
+		err = errno;
+		snprintf(err_str, 128, 
+			"RPNExpr.__setstate__() failed to allocate memory for \
+args buffer : %s",
+			strerror(err));
+		PyErr_SetString(PyExc_MemoryError, err_str);
+		return NULL;
+	}
+
+	if(rpn_expr_init(expr_self->rpn, state->stack_sz, state->argc) < 0)
+	{
+		snprintf(err_str, 256,
+			"RPNExpr.__setstate__() failed to init RPN expression \
+: %s",
+			expr_self->rpn->err_reason);
+		PyErr_SetString(PyExc_RuntimeError, err_str);
+		goto free_err;
+	}
+
+	/* restore stack & untokenize expression */
+	stack = (rpn_value_t*)(state+1);
+	memcpy(expr_self->rpn->stack, stack,
+		sizeof(rpn_value_t)*state->stack_sz);
+
+	toks.argc = state->argc;
+	toks.tokens_sz = state->token_sz;
+	toks.tokens = malloc(sizeof(rpn_token_t)*state->token_sz);
+	if(!toks.tokens)
+	{
+		snprintf(err_str, 128,
+			"RPNExpr.__setstate__() failed to allocate tokens \
+: %s",
+			strerror(errno));
+		PyErr_SetString(PyExc_RuntimeError, err_str);
+		goto close_err;
+	}
+	memcpy(toks.tokens, (rpn_token_t*)(stack+state->stack_sz),
+		sizeof(rpn_token_t)*state->token_sz);
+
+	for(i=0; i<toks.tokens_sz;i++)
+	{
+		// restore op pointers
+		toks.tokens[i].op = &(rpn_ops[toks.tokens[i].op_n]);
+	}
+	
+	
+	if(rpn_expr_untokenize(expr_self->rpn, &toks, 0) < 0)
+	{
+		snprintf(err_str, 256,
+			"RPNExpr.__setstate__() unable to untokenize : %s",
+			expr_self->rpn->err_reason);
+		PyErr_SetString(PyExc_RuntimeError, err_str);
+		goto close_err;
+	}
+
+	Py_RETURN_NONE;
+	close_err:
+		rpn_expr_close(expr_self->rpn);
+	free_err:
+		return NULL;
+}
+
 PyObject* rpnexpr_eval(PyObject* self, PyObject** argv, Py_ssize_t argc)
 {
 	PyRPNExpr_t *expr_self;

python_rpnexpr.h

@@ -62,9 +62,20 @@ typedef struct
 
 	/**@brief Array storing expression argument
 	 * @note As attribute of rpn_expr allowing malloc & free only once */
-	unsigned long *args;
+	rpn_value_t *args;
 } PyRPNExpr_t;
 
+/**@brief Organize PyRPNExpr_t state data
+ * @see rpnexpr_getstate
+ * @see rpnexpr_setstate */
+typedef struct
+{
+	size_t total_sz;
+	size_t argc;
+	unsigned char stack_sz;
+	size_t token_sz;
+} PyRPNExpr_state_t;
+
 /**@brief RpnExpr __new__ method
  * @param subtype Type of object being created (pyrpn.RPNExpr)
  * @param args positional arguments for subtype
@@ -87,6 +98,23 @@ int rpnexpr_init(PyObject *self, PyObject *args, PyObject *kwds);
  */
 void rpnexpr_del(PyObject *self);
 
+/**@brief RPNExpr __getstate__ method for pickling
+ * @param cls RPNExpr type object
+ * @param noargs Not an argument...
+ * @return A bytes Python instance suitable as argument for
+ * @ref rpnexpr_setstate
+ */
+PyObject* rpnexpr_getstate(PyObject *cls, PyObject *noargs);
+
+/**@brief RPNExpr __setstate__ method for pickling
+ * @param cls RPNExpr type object
+ * @param state Should by a bytes Python instance returned by @ref
+ * rpnexpr_getstate
+ * @return A bytes Python instance suitable as argument for
+ * @ref rpnexpr_setstate
+ */
+PyObject* rpnexpr_setstate(PyObject *cls, PyObject *state);
+
 /**@brief Eval an RPN expression given arguments and return the
  * value
  * @param self RPNExpr instance

test.c

@@ -219,6 +219,7 @@ int test_tokenization()
 				expr_orig, expr_untok);
 			return 1;
 		}
+
 		rpn_expr_close(&expr);
 		free(expr_untok);
 	}

tests/tests_rpn.py

@@ -20,6 +20,7 @@
 import sys
 import random
 import math
+import pickle
 
 import unittest
 from unittest import skip
@@ -144,7 +145,19 @@ class TestRpnCompile(unittest.TestCase):
             for j in range(256,4):
                 exprs = ' '.join(['+' for _ in range(i)])
                 expr = pyrpn.RPNExpr(exprs, j)
-            
+
+    def test_pickling(self):  
+        """ Pickle & unpickle tests """
+        argc = 4
+        for _ in range(512):
+            expr = pyrpn.RPNExpr(pyrpn.random_expr(2,10), argc)
+            pik = pickle.dumps(expr)
+            new_expr = pickle.loads(pik)
+            self.assertEqual(str(expr), str(new_expr))
+            args = [random.randint(0,IMAX) for _ in range(argc)]
+            self.assertEqual(expr.eval(*args), new_expr.eval(*args))
+
+
 
 class TestRpnEval(unittest.TestCase):