1 files changed, 53 insertions, 18 deletions

diff --git a/psycopg/adapter_qstring.c b/psycopg/adapter_qstring.c
index 2e3ab0a..1e256cf 100644
--- a/psycopg/adapter_qstring.c
+++ b/psycopg/adapter_qstring.c
@@ -36,28 +36,43 @@ static const char *default_encoding = "latin1";
 
 /* qstring_quote - do the quote process on plain and unicode strings */
 
+const char *
+_qstring_get_encoding(qstringObject *self)
+{
+    /* if the wrapped object is an unicode object we can encode it to match
+       conn->encoding but if the encoding is not specified we don't know what
+       to do and we raise an exception */
+    if (self->conn) {
+        return self->conn->codec;
+    }
+    else {
+        return self->encoding ? self->encoding : default_encoding;
+    }
+}
+
 static PyObject *
 qstring_quote(qstringObject *self)
 {
     PyObject *str = NULL;
     char *s, *buffer = NULL;
     Py_ssize_t len, qlen;
-    const char *encoding = default_encoding;
+    const char *encoding;
     PyObject *rv = NULL;
 
-    /* if the wrapped object is an unicode object we can encode it to match
-       conn->encoding but if the encoding is not specified we don't know what
-       to do and we raise an exception */
-    if (self->conn) {
-        encoding = self->conn->codec;
-    }
-
+    encoding = _qstring_get_encoding(self);
     Dprintf("qstring_quote: encoding to %s", encoding);
 
-    if (PyUnicode_Check(self->wrapped) && encoding) {
-        str = PyUnicode_AsEncodedString(self->wrapped, encoding, NULL);
-        Dprintf("qstring_quote: got encoded object at %p", str);
-        if (str == NULL) goto exit;
+    if (PyUnicode_Check(self->wrapped)) {
+        if (encoding) {
+            str = PyUnicode_AsEncodedString(self->wrapped, encoding, NULL);
+            Dprintf("qstring_quote: got encoded object at %p", str);
+            if (str == NULL) goto exit;
+        }
+        else {
+            PyErr_SetString(PyExc_TypeError,
+                "missing encoding to encode unicode object");
+            goto exit;
+        }
     }
 
 #if PY_MAJOR_VERSION < 3
@@ -150,13 +165,32 @@ qstring_conform(qstringObject *self, PyObject *args)
 static PyObject *
 qstring_get_encoding(qstringObject *self)
 {
-    const char *encoding = default_encoding;
+    const char *encoding;
+    encoding = _qstring_get_encoding(self);
+    return Text_FromUTF8(encoding);
+}
 
-    if (self->conn) {
-        encoding = self->conn->codec;
-    }
+static int
+qstring_set_encoding(qstringObject *self, PyObject *pyenc)
+{
+    int rv = -1;
+    const char *tmp;
+    char *cenc;
 
-    return Text_FromUTF8(encoding);
+    /* get a C copy of the encoding (which may come from unicode) */
+    Py_INCREF(pyenc);
+    if (!(pyenc = psycopg_ensure_bytes(pyenc))) { goto exit; }
+    if (!(tmp = Bytes_AsString(pyenc))) { goto exit; }
+    if (0 > psycopg_strdup(&cenc, tmp, 0)) { goto exit; }
+
+    Dprintf("qstring_set_encoding: encoding set to %s", cenc);
+    PyMem_Free((void *)self->encoding);
+    self->encoding = cenc;
+    rv = 0;
+
+exit:
+    Py_XDECREF(pyenc);
+    return rv;
 }
 
 /** the QuotedString object **/
@@ -183,7 +217,7 @@ static PyMethodDef qstringObject_methods[] = {
 static PyGetSetDef qstringObject_getsets[] = {
     { "encoding",
         (getter)qstring_get_encoding,
-        (setter)NULL,
+        (setter)qstring_set_encoding,
         "current encoding of the adapter" },
     {NULL}
 };
@@ -216,6 +250,7 @@ qstring_dealloc(PyObject* obj)
     Py_CLEAR(self->wrapped);
     Py_CLEAR(self->buffer);
     Py_CLEAR(self->conn);
+    PyMem_Free((void *)self->encoding);
 
     Dprintf("qstring_dealloc: deleted qstring object at %p, refcnt = "
         FORMAT_CODE_PY_SSIZE_T,