diff options
| -rw-r--r-- | psycopg/psycopg.h | 1 | ||||
| -rw-r--r-- | psycopg/python.h | 2 | ||||
| -rw-r--r-- | psycopg/utils.c | 26 |
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; +} + |