[pypy-commit] pypy default: cpyext: implement PyThreadState_GetDict()

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: 
Changeset: r52527:4a55cbff44eb
Date: 2012-02-15 21:24 +0100
http://bitbucket.org/pypy/pypy/changeset/4a55cbff44eb/

Log:	cpyext: implement PyThreadState_GetDict()

diff --git a/pypy/module/cpyext/include/pystate.h b/pypy/module/cpyext/include/pystate.h
--- a/pypy/module/cpyext/include/pystate.h
+++ b/pypy/module/cpyext/include/pystate.h
@@ -10,6 +10,7 @@
 
 typedef struct _ts {
     PyInterpreterState *interp;
+    PyObject *dict;  /* Stores per-thread state */
 } PyThreadState;
 
 #define Py_BEGIN_ALLOW_THREADS { \
diff --git a/pypy/module/cpyext/include/pythread.h b/pypy/module/cpyext/include/pythread.h
--- a/pypy/module/cpyext/include/pythread.h
+++ b/pypy/module/cpyext/include/pythread.h
@@ -1,6 +1,8 @@
 #ifndef Py_PYTHREAD_H
 #define Py_PYTHREAD_H
 
+#define WITH_THREAD
+
 typedef void *PyThread_type_lock;
 #define WAIT_LOCK	1
 #define NOWAIT_LOCK	0
diff --git a/pypy/module/cpyext/pystate.py b/pypy/module/cpyext/pystate.py
--- a/pypy/module/cpyext/pystate.py
+++ b/pypy/module/cpyext/pystate.py
@@ -1,12 +1,19 @@
 from pypy.module.cpyext.api import (
     cpython_api, generic_cpy_call, CANNOT_FAIL, CConfig, cpython_struct)
+from pypy.module.cpyext.pyobject import PyObject, Py_DecRef, make_ref
 from pypy.rpython.lltypesystem import rffi, lltype
 
 PyInterpreterStateStruct = lltype.ForwardReference()
 PyInterpreterState = lltype.Ptr(PyInterpreterStateStruct)
 cpython_struct(
-    "PyInterpreterState", [('next', PyInterpreterState)], PyInterpreterStateStruct)
-PyThreadState = lltype.Ptr(cpython_struct("PyThreadState", [('interp', PyInterpreterState)]))
+    "PyInterpreterState",
+    [('next', PyInterpreterState)],
+    PyInterpreterStateStruct)
+PyThreadState = lltype.Ptr(cpython_struct(
+    "PyThreadState", 
+    [('interp', PyInterpreterState),
+     ('dict', PyObject),
+     ]))
 
 @cpython_api([], PyThreadState, error=CANNOT_FAIL)
 def PyEval_SaveThread(space):
@@ -38,41 +45,48 @@
     return 1
 
 # XXX: might be generally useful
-def encapsulator(T, flavor='raw'):
+def encapsulator(T, flavor='raw', dealloc=None):
     class MemoryCapsule(object):
-        def __init__(self, alloc=True):
-            if alloc:
+        def __init__(self, space):
+            self.space = space
+            if space is not None:
                 self.memory = lltype.malloc(T, flavor=flavor)
             else:
                 self.memory = lltype.nullptr(T)
         def __del__(self):
             if self.memory:
+                if dealloc and self.space:
+                    dealloc(self.memory, self.space)
                 lltype.free(self.memory, flavor=flavor)
     return MemoryCapsule
 
-ThreadStateCapsule = encapsulator(PyThreadState.TO)
+def ThreadState_dealloc(ts, space):
+    Py_DecRef(space, ts.c_dict)
+ThreadStateCapsule = encapsulator(PyThreadState.TO,
+                                  dealloc=ThreadState_dealloc)
 
 from pypy.interpreter.executioncontext import ExecutionContext
-ExecutionContext.cpyext_threadstate = ThreadStateCapsule(alloc=False)
+ExecutionContext.cpyext_threadstate = ThreadStateCapsule(None)
 
 class InterpreterState(object):
     def __init__(self, space):
         self.interpreter_state = lltype.malloc(
             PyInterpreterState.TO, flavor='raw', zero=True, immortal=True)
 
-    def new_thread_state(self):
-        capsule = ThreadStateCapsule()
+    def new_thread_state(self, space):
+        capsule = ThreadStateCapsule(space)
         ts = capsule.memory
         ts.c_interp = self.interpreter_state
+        ts.c_dict = make_ref(space, space.newdict())
         return capsule
 
     def get_thread_state(self, space):
         ec = space.getexecutioncontext()
-        return self._get_thread_state(ec).memory
+        return self._get_thread_state(space, ec).memory
 
-    def _get_thread_state(self, ec):
+    def _get_thread_state(self, space, ec):
         if ec.cpyext_threadstate.memory == lltype.nullptr(PyThreadState.TO):
-            ec.cpyext_threadstate = self.new_thread_state()
+            ec.cpyext_threadstate = self.new_thread_state(space)
 
         return ec.cpyext_threadstate
 
@@ -81,6 +95,11 @@
     state = space.fromcache(InterpreterState)
     return state.get_thread_state(space)
 
+ at cpython_api([], PyObject, error=CANNOT_FAIL)
+def PyThreadState_GetDict(space):
+    state = space.fromcache(InterpreterState)
+    return state.get_thread_state(space).c_dict
+
 @cpython_api([PyThreadState], PyThreadState, error=CANNOT_FAIL)
 def PyThreadState_Swap(space, tstate):
     """Swap the current thread state with the thread state given by the argument
diff --git a/pypy/module/cpyext/test/test_pystate.py b/pypy/module/cpyext/test/test_pystate.py
--- a/pypy/module/cpyext/test/test_pystate.py
+++ b/pypy/module/cpyext/test/test_pystate.py
@@ -2,6 +2,7 @@
 from pypy.module.cpyext.test.test_api import BaseApiTest
 from pypy.rpython.lltypesystem.lltype import nullptr
 from pypy.module.cpyext.pystate import PyInterpreterState, PyThreadState
+from pypy.module.cpyext.pyobject import from_ref
 
 class AppTestThreads(AppTestCpythonExtensionBase):
     def test_allow_threads(self):
@@ -49,3 +50,10 @@
 
         api.PyEval_AcquireThread(tstate)
         api.PyEval_ReleaseThread(tstate)
+
+    def test_threadstate_dict(self, space, api):
+        ts = api.PyThreadState_Get()
+        ref = ts.c_dict
+        assert ref == api.PyThreadState_GetDict()
+        w_obj = from_ref(space, ref)
+        assert space.isinstance_w(w_obj, space.w_dict)