From bccbcf42d0c33f1a94cad26647b220bd541179e7 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 16 Feb 2011 01:23:53 +0000
Subject: [PATCH] Added adaptation for objects supporting the new-style buffer
 interface

Supporting this interface is required to adapt memoryview on Python 2.7 as they
don't support the old style. But because the old style is long deprecated it
makes sense to start supporting the new one.
---
 psycopg/adapter_binary.c | 98 ++++++++++++++++++++++++----------------
 1 file changed, 59 insertions(+), 39 deletions(-)

diff --git a/psycopg/adapter_binary.c b/psycopg/adapter_binary.c
index 31b25026..00bdb9ee 100644
--- a/psycopg/adapter_binary.c
+++ b/psycopg/adapter_binary.c
@@ -47,57 +47,79 @@ binary_escape(unsigned char *from, size_t from_length,
         return PQescapeBytea(from, from_length, to_length);
 }
 
+#define HAS_BUFFER (PY_MAJOR_VERSION < 3)
+#define HAS_MEMORYVIEW (PY_MAJOR_VERSION > 2 || PY_MINOR_VERSION >= 6)
+
 /* binary_quote - do the quote process on plain and unicode strings */
 
 static PyObject *
 binary_quote(binaryObject *self)
 {
-    char *to;
-    const char *buffer;
+    char *to = NULL;
+    const char *buffer = NULL;
     Py_ssize_t buffer_len;
     size_t len = 0;
+    PyObject *rv = NULL;
+#if HAS_MEMORYVIEW
+    Py_buffer view;
+    int got_view = 0;
+#endif
 
     /* if we got a plain string or a buffer we escape it and save the buffer */
-    if (Bytes_Check(self->wrapped)
-#if PY_MAJOR_VERSION < 3
-        || PyBuffer_Check(self->wrapped)
-#endif
-#if PY_MAJOR_VERSION >= 3 || PY_MINOR_VERSION >= 6
-        || PyByteArray_Check(self->wrapped)
-#endif
-#if PY_MAJOR_VERSION >= 3 || PY_MINOR_VERSION >= 7
-        || PyMemoryView_Check(self->wrapped)
-#endif
-        ) {
-        /* escape and build quoted buffer */
-        if (PyObject_AsReadBuffer(self->wrapped, (const void **)&buffer,
-                                  &buffer_len) < 0)
-            return NULL;
 
-        to = (char *)binary_escape((unsigned char*)buffer, (size_t) buffer_len,
-            &len, self->conn ? ((connectionObject*)self->conn)->pgconn : NULL);
-        if (to == NULL) {
-            PyErr_NoMemory();
-            return NULL;
+#if HAS_MEMORYVIEW
+    if (PyObject_CheckBuffer(self->wrapped)) {
+        if (0 > PyObject_GetBuffer(self->wrapped, &view, PyBUF_CONTIG_RO)) {
+            goto exit;
         }
-
-        if (len > 0)
-            self->buffer = Bytes_FromFormat(
-                (self->conn && ((connectionObject*)self->conn)->equote)
-                    ? "E'%s'::bytea" : "'%s'::bytea" , to);
-        else
-            self->buffer = Bytes_FromString("''::bytea");
-
-        PQfreemem(to);
+        got_view = 1;
+        buffer = (const char *)(view.buf);
+        buffer_len = view.len;
     }
+#endif
+
+#if HAS_BUFFER
+    if (!buffer && (Bytes_Check(self->wrapped) || PyBuffer_Check(self->wrapped))) {
+        if (PyObject_AsReadBuffer(self->wrapped, (const void **)&buffer,
+                                  &buffer_len) < 0) {
+            goto exit;
+        }
+    }
+#endif
+
+    if (!buffer) {
+        goto exit;
+    }
+
+    /* escape and build quoted buffer */
+
+    to = (char *)binary_escape((unsigned char*)buffer, (size_t) buffer_len,
+        &len, self->conn ? ((connectionObject*)self->conn)->pgconn : NULL);
+    if (to == NULL) {
+        PyErr_NoMemory();
+        goto exit;
+    }
+
+    if (len > 0)
+        rv = Bytes_FromFormat(
+            (self->conn && ((connectionObject*)self->conn)->equote)
+                ? "E'%s'::bytea" : "'%s'::bytea" , to);
+    else
+        rv = Bytes_FromString("''::bytea");
+
+exit:
+    if (to) { PQfreemem(to); }
+#if HAS_MEMORYVIEW
+    if (got_view) { PyBuffer_Release(&view); }
+#endif
 
     /* if the wrapped object is not bytes or a buffer, this is an error */
-    else {
-        PyErr_SetString(PyExc_TypeError, "can't escape non-bytes object");
-        return NULL;
+    if (!rv && !PyErr_Occurred()) {
+        PyErr_Format(PyExc_TypeError, "can't escape %s to binary",
+            Py_TYPE(self->wrapped)->tp_name);
     }
 
-    return self->buffer;
+    return rv;
 }
 
 /* binary_str, binary_getquoted - return result of quoting */
@@ -106,11 +128,9 @@ static PyObject *
 binary_getquoted(binaryObject *self, PyObject *args)
 {
     if (self->buffer == NULL) {
-        if (!(binary_quote(self))) {
-            return NULL;
-        }
+        self->buffer = binary_quote(self);
     }
-    Py_INCREF(self->buffer);
+    Py_XINCREF(self->buffer);
     return self->buffer;
 }