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

3 files changed, 29 insertions, 0 deletions

diff --git a/psycopg/psycopg.h b/psycopg/psycopg.h
index 00cb98a..57a73d0 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 452c734..ee2dee7 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 f2a0ab0..e5b221f 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;
+}
+