From b4685bba4afff9eedfa0b03c1de815cfe738151f Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Thu, 23 Dec 2010 19:10:07 +0100
Subject: [PATCH] Added utility function to get bytes from a str/unicode.

---
 psycopg/psycopg.h |  1 +
 psycopg/python.h  |  2 ++
 psycopg/utils.c   | 26 ++++++++++++++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/psycopg/psycopg.h b/psycopg/psycopg.h
index 00cb98ad..57a73d02 100644
--- a/psycopg/psycopg.h
+++ b/psycopg/psycopg.h
@@ -121,6 +121,7 @@ HIDDEN char *psycopg_escape_string(PyObject *conn,
               const char *from, Py_ssize_t len, char *to, Py_ssize_t *tolen);
 
 HIDDEN char *psycopg_strdup(const char *from, Py_ssize_t len);
+HIDDEN PyObject * psycopg_ensure_bytes(PyObject *obj);
 
 /* Exceptions docstrings */
 #define Error_doc \
diff --git a/psycopg/python.h b/psycopg/python.h
index 452c7341..ee2dee76 100644
--- a/psycopg/python.h
+++ b/psycopg/python.h
@@ -106,6 +106,7 @@
 /* XXX BytesType -> Bytes_Type */
 #define BytesType PyString_Type
 #define Bytes_Check PyString_Check
+#define Bytes_CheckExact PyString_CheckExact
 #define Bytes_AS_STRING PyString_AS_STRING
 #define Bytes_GET_SIZE PyString_GET_SIZE
 #define Bytes_Size PyString_Size
@@ -120,6 +121,7 @@
 
 #define BytesType PyBytes_Type
 #define Bytes_Check PyBytes_Check
+#define Bytes_CheckExact PyBytes_CheckExact
 #define Bytes_AS_STRING PyBytes_AS_STRING
 #define Bytes_GET_SIZE PyBytes_GET_SIZE
 #define Bytes_Size PyBytes_Size
diff --git a/psycopg/utils.c b/psycopg/utils.c
index f2a0ab09..e5b221f9 100644
--- a/psycopg/utils.c
+++ b/psycopg/utils.c
@@ -92,3 +92,29 @@ psycopg_strdup(const char *from, Py_ssize_t len)
     return rv;
 }
 
+/* Ensure a Python object is a bytes string.
+ *
+ * Useful when a char * is required out of it.
+ *
+ * Return a new reference. NULL on error.
+ */
+PyObject *
+psycopg_ensure_bytes(PyObject *obj)
+{
+    PyObject *rv = NULL;
+
+    if (PyUnicode_CheckExact(obj)) {
+        rv = PyUnicode_AsUTF8String(obj);
+    }
+    else if (Bytes_CheckExact(obj)) {
+        Py_INCREF(obj);
+        rv = obj;
+    }
+    else {
+        PyErr_Format(PyExc_TypeError, "I'm not into ensuring %s as bytes",
+            obj ? Py_TYPE(obj)->tp_name : "NULL");
+    }
+
+    return rv;
+}
+