Function psycopg_ensure_bytes converted in a "filter" stealing a ref.

4 files changed, 21 insertions, 21 deletions

diff --git a/psycopg/adapter_datetime.c b/psycopg/adapter_datetime.c
index 56be0c9..ddcd089 100644
--- a/psycopg/adapter_datetime.c
+++ b/psycopg/adapter_datetime.c
@@ -59,7 +59,6 @@ _pydatetime_string_date_time(pydatetimeObject *self)
 {
     PyObject *rv = NULL;
     PyObject *iso = NULL;
-    PyObject *biso = NULL;
     PyObject *tz;
 
     /* Select the right PG type to cast into. */
@@ -79,22 +78,17 @@ _pydatetime_string_date_time(pydatetimeObject *self)
         break;
     }
 
-    if (!(iso = PyObject_CallMethod(self->wrapped, "isoformat", NULL))) {
+    if (!(iso = psycopg_ensure_bytes(
+            PyObject_CallMethod(self->wrapped, "isoformat", NULL)))) {
         goto error;
     }
 
-    if (!(biso = psycopg_ensure_bytes(iso))) {
-        goto error;
-    }
-
-    rv = Bytes_FromFormat(fmt, Bytes_AsString(biso));
+    rv = Bytes_FromFormat(fmt, Bytes_AsString(iso));
 
-    Py_DECREF(biso);
     Py_DECREF(iso);
     return rv;
 
 error:
-    Py_XDECREF(biso);
     Py_XDECREF(iso);
     return rv;
 }
diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index ee4bebc..6d281b3 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -239,7 +239,6 @@ conn_encoding_to_codec(const char *enc)
     char *tmp;
     Py_ssize_t size;
     PyObject *pyenc = NULL;
-    PyObject *pybenc = NULL;
     char *rv = NULL;
 
     /* Find the Py codec name from the PG encoding */
@@ -250,11 +249,12 @@ conn_encoding_to_codec(const char *enc)
     }
 
     /* Convert the codec in a bytes string to extract the c string. */
-    if (!(pybenc = psycopg_ensure_bytes(pyenc))) {
+    Py_INCREF(pyenc);
+    if (!(pyenc = psycopg_ensure_bytes(pyenc))) {
         goto exit;
     }
 
-    if (-1 == Bytes_AsStringAndSize(pybenc, &tmp, &size)) {
+    if (-1 == Bytes_AsStringAndSize(pyenc, &tmp, &size)) {
         goto exit;
     }
 
@@ -262,8 +262,7 @@ conn_encoding_to_codec(const char *enc)
     rv = psycopg_strdup(tmp, size);
 
 exit:
-    /* pyenc is borrowed: no decref. */
-    Py_XDECREF(pybenc);
+    Py_XDECREF(pyenc);
     return rv;
 }
 
diff --git a/psycopg/typecast.c b/psycopg/typecast.c
index b59edb0..cc155f4 100644
--- a/psycopg/typecast.c
+++ b/psycopg/typecast.c
@@ -430,17 +430,17 @@ static PyObject *
 typecast_repr(PyObject *self)
 {
     PyObject *name = ((typecastObject *)self)->name;
-    PyObject *bname;
     PyObject *rv;
 
-    if (!(bname = psycopg_ensure_bytes(name))) {
+    Py_INCREF(name);
+    if (!(name = psycopg_ensure_bytes(name))) {
         return NULL;
     }
 
     rv = PyString_FromFormat("<%s '%s' at %p>",
         Py_TYPE(self)->tp_name, PyBytes_AS_STRING(name), self);
 
-    Py_DECREF(bname);
+    Py_DECREF(name);
     return rv;
 }
 
diff --git a/psycopg/utils.c b/psycopg/utils.c
index 16b9249..539081b 100644
--- a/psycopg/utils.c
+++ b/psycopg/utils.c
@@ -96,23 +96,30 @@ psycopg_strdup(const char *from, Py_ssize_t len)
  *
  * Useful when a char * is required out of it.
  *
- * Return a new reference. NULL on error.
+ * The function is ref neutral: steals a ref from obj and adds one to the
+ * return value. This also means that you shouldn't call the function on a
+ * borrowed ref, if having the object unallocated is not what you want.
+ *
+ * It is safe to call the function on NULL.
  */
 PyObject *
 psycopg_ensure_bytes(PyObject *obj)
 {
     PyObject *rv = NULL;
+    if (!obj) { return NULL; }
 
     if (PyUnicode_CheckExact(obj)) {
         rv = PyUnicode_AsUTF8String(obj);
+        Py_DECREF(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");
+        PyErr_Format(PyExc_TypeError,
+            "Expected bytes or unicode string, got %s instead",
+            Py_TYPE(obj)->tp_name);
+        Py_DECREF(obj);  /* steal the ref anyway */
     }
 
     return rv;