From a4cc1c971d63059ffd0cfb2b1dd15ce57f4f61d6 Mon Sep 17 00:00:00 2001
From: Alexey Borzenkov <snaury@gmail.com>
Date: Fri, 11 May 2012 02:41:53 +0400
Subject: [PATCH] Add support for instance dictionaries

---
 greenlet.c             | 36 +++++++++++++++++++++++++++++++++++-
 greenlet.h             |  1 +
 tests/test_greenlet.py | 17 +++++++++++++++++
 3 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/greenlet.c b/greenlet.c
index 68caa4ef..addfc9ba 100644
--- a/greenlet.c
+++ b/greenlet.c
@@ -728,6 +728,7 @@ green_traverse(PyGreenlet *self, visitproc visit, void *arg)
 	Py_VISIT(self->exc_type);
 	Py_VISIT(self->exc_value);
 	Py_VISIT(self->exc_traceback);
+	Py_VISIT(self->dict);
 	return 0;
 }
 
@@ -802,6 +803,7 @@ static void green_dealloc(PyGreenlet* self)
 	Py_CLEAR(self->exc_type);
 	Py_CLEAR(self->exc_value);
 	Py_CLEAR(self->exc_traceback);
+	Py_CLEAR(self->dict);
 	if (self->weakreflist != NULL)
 		PyObject_ClearWeakRefs((PyObject *) self);
 	Py_CLEAR(self->run_info);
@@ -978,6 +980,36 @@ static int green_bool(PyGreenlet* self)
 	return PyGreenlet_ACTIVE(self);
 }
 
+static PyObject* green_getdict(PyGreenlet* self, void* c)
+{
+	if (self->dict == NULL) {
+		self->dict = PyDict_New();
+		if (self->dict == NULL)
+			return NULL;
+	}
+	Py_INCREF(self->dict);
+	return self->dict;
+}
+
+static int green_setdict(PyGreenlet* self, PyObject* val, void* c)
+{
+	PyObject* tmp;
+
+	if (val == NULL) {
+		PyErr_SetString(PyExc_TypeError, "__dict__ may not be deleted");
+		return -1;
+	}
+	if (!PyDict_Check(val)) {
+		PyErr_SetString(PyExc_TypeError, "__dict__ must be a dictionary");
+		return -1;
+	}
+	tmp = self->dict;
+	Py_INCREF(val);
+	self->dict = val;
+	Py_XDECREF(tmp);
+	return 0;
+}
+
 static PyObject* green_getdead(PyGreenlet* self, void* c)
 {
 	PyObject* res;
@@ -1156,6 +1188,8 @@ static PyMethodDef green_methods[] = {
 };
 
 static PyGetSetDef green_getsets[] = {
+	{"__dict__", (getter)green_getdict,
+		     (setter)green_setdict, /*XXX*/ NULL},
 	{"run",    (getter)green_getrun,
 		   (setter)green_setrun, /*XXX*/ NULL},
 	{"parent", (getter)green_getparent,
@@ -1229,7 +1263,7 @@ PyTypeObject PyGreenlet_Type = {
 	0,					/* tp_dict */
 	0,					/* tp_descr_get */
 	0,					/* tp_descr_set */
-	0,					/* tp_dictoffset */
+	offsetof(PyGreenlet, dict),		/* tp_dictoffset */
 	(initproc)green_init,			/* tp_init */
 	GREENLET_tp_alloc,			/* tp_alloc */
 	green_new,				/* tp_new */
diff --git a/greenlet.h b/greenlet.h
index 1d0d28f8..7370f6aa 100644
--- a/greenlet.h
+++ b/greenlet.h
@@ -28,6 +28,7 @@ typedef struct _greenlet {
 	PyObject* exc_type;
 	PyObject* exc_value;
 	PyObject* exc_traceback;
+	PyObject* dict;
 } PyGreenlet;
 
 #define PyGreenlet_Check(op)      PyObject_TypeCheck(op, &PyGreenlet_Type)
diff --git a/tests/test_greenlet.py b/tests/test_greenlet.py
index 37f678cd..33aad0f8 100644
--- a/tests/test_greenlet.py
+++ b/tests/test_greenlet.py
@@ -245,3 +245,20 @@ def h():
             self.assertEqual(sys.exc_info(), (None, None, None))
 
         greenlet(f).switch()
+
+    def test_instance_dict(self):
+        def f():
+            greenlet.getcurrent().test = 42
+        def deldict(g):
+            del g.__dict__
+        def setdict(g, value):
+            g.__dict__ = value
+        g = greenlet(f)
+        self.assertEqual(g.__dict__, {})
+        g.switch()
+        self.assertEqual(g.test, 42)
+        self.assertEqual(g.__dict__, {'test': 42})
+        g.__dict__ = g.__dict__
+        self.assertEqual(g.__dict__, {'test': 42})
+        self.assertRaises(TypeError, deldict, g)
+        self.assertRaises(TypeError, setdict, g, 42)