[pypy-commit] pypy rw-PyString_AS_STRING: allow rw access to the pointer returned from PyString_AS_STRING(obj) if obj is newly created

[mattip]

Author: Matti Picus <matti.picus at gmail.com>
Branch: rw-PyString_AS_STRING
Changeset: r85155:5b2e7547c432
Date: 2016-06-14 12:23 +0300
http://bitbucket.org/pypy/pypy/changeset/5b2e7547c432/

Log:	allow rw access to the pointer returned from PyString_AS_STRING(obj)
	if obj is newly created

diff --git a/pypy/module/cpyext/bytesobject.py b/pypy/module/cpyext/bytesobject.py
--- a/pypy/module/cpyext/bytesobject.py
+++ b/pypy/module/cpyext/bytesobject.py
@@ -6,7 +6,8 @@
 from pypy.module.cpyext.pyerrors import PyErr_BadArgument
 from pypy.module.cpyext.pyobject import (
     PyObject, PyObjectP, Py_DecRef, make_ref, from_ref, track_reference,
-    make_typedescr, get_typedescr, as_pyobj, Py_IncRef, get_w_obj_and_decref)
+    make_typedescr, get_typedescr, as_pyobj, Py_IncRef, get_w_obj_and_decref,
+    pyobj_has_w_obj)
 
 ##
 ## Implementation of PyStringObject
@@ -41,6 +42,9 @@
 ##   the pypy string is created, and added to the global map of tracked
 ##   objects.  The buffer is then supposed to be immutable.
 ##
+##-  A buffer obtained from PyString_AS_STRING() could be mutable iff
+##   there is no corresponding pypy object for the string
+##
 ## - _PyString_Resize() works only on not-yet-pypy'd strings, and returns a
 ##   similar object.
 ##
@@ -139,6 +143,9 @@
 
 @cpython_api([PyObject], rffi.CCHARP, error=0)
 def PyString_AsString(space, ref):
+    return _PyString_AsString(space, ref)
+
+def _PyString_AsString(space, ref):
     if from_ref(space, rffi.cast(PyObject, ref.c_ob_type)) is space.w_str:
         pass    # typecheck returned "ok" without forcing 'ref' at all
     elif not PyString_Check(space, ref):   # otherwise, use the alternate way
@@ -158,6 +165,19 @@
         ref_str.c_buffer = rffi.str2charp(s)
     return ref_str.c_buffer
 
+ at cpython_api([rffi.VOIDP], rffi.CCHARP, error=0)
+def PyString_AS_STRING(space, void_ref):
+    ref = rffi.cast(PyObject, void_ref)
+    # if no w_str is associated with this ref,
+    # return the c-level ptr as RW
+    if not pyobj_has_w_obj(ref):
+        py_str = rffi.cast(PyStringObject, ref)
+        if not py_str.c_buffer:
+            py_str.c_buffer = lltype.malloc(rffi.CCHARP.TO, py_str.c_ob_size + 1,
+                                        flavor='raw', zero=True)
+        return py_str.c_buffer
+    return _PyString_AsString(space, ref)
+
 @cpython_api([PyObject, rffi.CCHARPP, rffi.CArrayPtr(Py_ssize_t)], rffi.INT_real, error=-1)
 def PyString_AsStringAndSize(space, ref, buffer, length):
     if not PyString_Check(space, ref):
diff --git a/pypy/module/cpyext/include/stringobject.h b/pypy/module/cpyext/include/stringobject.h
--- a/pypy/module/cpyext/include/stringobject.h
+++ b/pypy/module/cpyext/include/stringobject.h
@@ -10,7 +10,6 @@
 #include <stdarg.h>
 
 #define PyString_GET_SIZE(op) PyString_Size((PyObject*)(op))
-#define PyString_AS_STRING(op) PyString_AsString((PyObject*)(op))
 /*
 Type PyStringObject represents a character string.  An extra zero byte is
 reserved at the end to ensure it is zero-terminated, but a size is
diff --git a/pypy/module/cpyext/test/test_bytesobject.py b/pypy/module/cpyext/test/test_bytesobject.py
--- a/pypy/module/cpyext/test/test_bytesobject.py
+++ b/pypy/module/cpyext/test/test_bytesobject.py
@@ -110,14 +110,23 @@
                 obj = (PyStringObject*)type->tp_alloc(type, 10);
                 if (PyString_GET_SIZE(obj) != 10)
                     return PyLong_FromLong(PyString_GET_SIZE(obj));
-                /* cannot work, there is only RO access
-                memcpy(PyString_AS_STRING(obj), "works", 6); */
+                /* cannot work, there is only RO access */
+                memcpy(PyString_AS_STRING(obj), "works", 6);
                 Py_INCREF(obj);
                 return (PyObject*)obj;
              """),
+            ('alloc_rw', "METH_NOARGS",
+             '''
+                PyObject *obj = _PyObject_NewVar(&PyString_Type, 10);
+                char * buf = PyString_AS_STRING(obj);
+                memcpy(PyString_AS_STRING(obj), "works", 6);
+                return (PyObject*)obj;
+             '''),
             ])
+        s = module.alloc_rw()
+        assert s == 'works' + '\x00' * 5
         s = module.tpalloc()
-        assert s == '\x00' * 10
+        assert s == 'works' + '\x00' * 5
 
     def test_AsString(self):
         module = self.import_extension('foo', [