lib/ldb: Implement a bytes derived object for attributes py2/py3

ldb attributes are either bytes (py3) or str (py2)

Some places in the code do str(res[0]['attribute'][0])
which results in
   'result' (py2)
  b'result' (py3)

or more commonly the attribute is used to construct a string e.g.
   "blah=" + res[0]['attribute'][0] + ",foo,bar=...."

giving
   "blah=result,foo,bar=...." (py2)
and very unhelpfully
   "blah=b'result',foo,bar=...." (py3)

lots of code already constructs various strings for passing to other
api using the above. To avoid many excessive
    res[0]['attribute'][0].decode('utf8')

code like 'res[0]['attribute'][0]'

will now return LdbBytes (a new object subclassing 'bytes') in py3
instead of bytes. This object has a custom '__str__' method which
attempts to return a string decoded to uft8. In Py2 this will behave as
it did previously (this is the safer option at the moment)

Signed-off-by: Noel Power <noel.power@suse.com>
Reviewed-by: Andreas Schneider <asn@samba.org>
Reviewed-by: Douglas Bagnall <douglas.bagnall@catalyst.net.nz>

1 files changed, 43 insertions, 2 deletions

diff --git a/lib/ldb/pyldb.c b/lib/ldb/pyldb.c
index 110ec8e60ad..66bc2021fb4 100644
--- a/lib/ldb/pyldb.c
+++ b/lib/ldb/pyldb.c
@@ -79,6 +79,7 @@ static struct ldb_message_element *PyObject_AsMessageElement(
 						      PyObject *set_obj,
 						      unsigned int flags,
 						      const char *attr_name);
+static PyTypeObject PyLdbBytesType;
 
 #if PY_MAJOR_VERSION >= 3
 #define PyStr_Check PyUnicode_Check
@@ -89,6 +90,16 @@ static struct ldb_message_element *PyObject_AsMessageElement(
 #define PyStr_AsUTF8 PyUnicode_AsUTF8
 #define PyStr_AsUTF8AndSize PyUnicode_AsUTF8AndSize
 #define PyInt_FromLong PyLong_FromLong
+
+static PyObject *PyLdbBytes_FromStringAndSize(const char *msg, int size)
+{
+	PyObject* result = NULL;
+	PyObject* args = NULL;
+	args = Py_BuildValue("(y#)", msg, size);
+	result = PyLdbBytesType.tp_new(&PyLdbBytesType, args, NULL);
+	Py_DECREF(args);
+	return result;
+}
 #else
 #define PyStr_Check PyString_Check
 #define PyStr_FromString PyString_FromString
@@ -96,6 +107,7 @@ static struct ldb_message_element *PyObject_AsMessageElement(
 #define PyStr_FromFormat PyString_FromFormat
 #define PyStr_FromFormatV PyString_FromFormatV
 #define PyStr_AsUTF8 PyString_AsString
+#define PyLdbBytes_FromStringAndSize PyString_FromStringAndSize
 
 const char *PyStr_AsUTF8AndSize(PyObject *pystr, Py_ssize_t *sizeptr);
 const char *
@@ -270,10 +282,34 @@ static void PyErr_SetLdbError(PyObject *error, int ret, struct ldb_context *ldb_
 			Py_BuildValue(discard_const_p(char, "(i,s)"), ret,
 				      ldb_ctx == NULL?ldb_strerror(ret):ldb_errstring(ldb_ctx)));
 }
+static PyObject *py_ldb_bytes_str(PyBytesObject *self)
+{
+	char *msg = NULL;
+	Py_ssize_t size;
+	int result = 0;
+	if (!PyBytes_Check(self)) {
+		PyErr_Format(PyExc_TypeError,"Unexpected type");
+		return NULL;
+	}
+	result = PyBytes_AsStringAndSize((PyObject *)self, &msg, &size);
+	if (result != 0) {
+		PyErr_Format(PyExc_TypeError, "Failed to extract bytes");
+		return NULL;
+	}
+	return PyUnicode_FromStringAndSize(msg, size);
+}
+
+static PyTypeObject PyLdbBytesType = {
+	PyVarObject_HEAD_INIT(NULL, 0)
+	.tp_name = "ldb.bytes",
+	.tp_doc = "str/bytes (with custom str)",
+        .tp_str = (reprfunc)py_ldb_bytes_str,
+	.tp_flags = Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE,
+};
 
 static PyObject *PyObject_FromLdbValue(const struct ldb_val *val)
 {
-	return PyBytes_FromStringAndSize((const char *)val->data, val->length);
+	return PyLdbBytes_FromStringAndSize((const char *)val->data, val->length);
 }
 
 static PyObject *PyStr_FromLdbValue(const struct ldb_val *val)
@@ -2990,7 +3026,7 @@ static PyObject *py_ldb_msg_element_find(PyLdbMessageElementObject *self, Py_ssi
 		PyErr_SetString(PyExc_IndexError, "Out of range");
 		return NULL;
 	}
-	return PyBytes_FromStringAndSize((char *)el->values[idx].data, el->values[idx].length);
+	return PyLdbBytes_FromStringAndSize((char *)el->values[idx].data, el->values[idx].length);
 }
 
 static PySequenceMethods py_ldb_msg_element_seq = {
@@ -4123,6 +4159,11 @@ static PyObject* module_init(void)
 {
 	PyObject *m;
 
+	PyLdbBytesType.tp_base = &PyBytes_Type;
+	if (PyType_Ready(&PyLdbBytesType) < 0) {
+		return NULL;
+	}
+
 	if (PyType_Ready(&PyLdbDn) < 0)
 		return NULL;