Added utility function to get bytes from a str/unicode.

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 \

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

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;
+}
+