#4759: allow None as first argument of bytearray.translate(), for consistency with bytes.translate().

Also fix segfault for bytearray.translate(x, None) -- will backport this part to 3.0 and 2.6.

1 files changed, 25 insertions, 16 deletions

diff --git a/Objects/bytearrayobject.c b/Objects/bytearrayobject.c
index 1620652..1c524c1 100644
--- a/Objects/bytearrayobject.c
+++ b/Objects/bytearrayobject.c
@@ -1371,28 +1371,32 @@ bytes_translate(PyByteArrayObject *self, PyObject *args)
     PyObject *input_obj = (PyObject*)self;
     const char *output_start;
     Py_ssize_t inlen;
-    PyObject *result;
+    PyObject *result = NULL;
     int trans_table[256];
-    PyObject *tableobj, *delobj = NULL;
+    PyObject *tableobj = NULL, *delobj = NULL;
     Py_buffer vtable, vdel;
 
     if (!PyArg_UnpackTuple(args, "translate", 1, 2,
                            &tableobj, &delobj))
           return NULL;
 
-    if (_getbuffer(tableobj, &vtable) < 0)
+    if (tableobj == Py_None) {
+        table = NULL;
+        tableobj = NULL;
+    } else if (_getbuffer(tableobj, &vtable) < 0) {
         return NULL;
-
-    if (vtable.len != 256) {
-        PyErr_SetString(PyExc_ValueError,
-                        "translation table must be 256 characters long");
-        result = NULL;
-        goto done;
+    } else {
+        if (vtable.len != 256) {
+            PyErr_SetString(PyExc_ValueError,
+                            "translation table must be 256 characters long");
+            goto done;
+        }
+        table = (const char*)vtable.buf;
     }
 
     if (delobj != NULL) {
         if (_getbuffer(delobj, &vdel) < 0) {
-            result = NULL;
+            delobj = NULL;  /* don't try to release vdel buffer on exit */
             goto done;
         }
     }
@@ -1401,7 +1405,6 @@ bytes_translate(PyByteArrayObject *self, PyObject *args)
         vdel.len = 0;
     }
 
-    table = (const char *)vtable.buf;
     inlen = PyByteArray_GET_SIZE(input_obj);
     result = PyByteArray_FromStringAndSize((char *)NULL, inlen);
     if (result == NULL)
@@ -1409,7 +1412,7 @@ bytes_translate(PyByteArrayObject *self, PyObject *args)
     output_start = output = PyByteArray_AsString(result);
     input = PyByteArray_AS_STRING(input_obj);
 
-    if (vdel.len == 0) {
+    if (vdel.len == 0 && table != NULL) {
         /* If no deletions are required, use faster code */
         for (i = inlen; --i >= 0; ) {
             c = Py_CHARMASK(*input++);
@@ -1417,9 +1420,14 @@ bytes_translate(PyByteArrayObject *self, PyObject *args)
         }
         goto done;
     }
-    
-    for (i = 0; i < 256; i++)
-        trans_table[i] = Py_CHARMASK(table[i]);
+
+    if (table == NULL) {
+        for (i = 0; i < 256; i++)
+            trans_table[i] = Py_CHARMASK(i);
+    } else {
+        for (i = 0; i < 256; i++)
+            trans_table[i] = Py_CHARMASK(table[i]);
+    }
 
     for (i = 0; i < vdel.len; i++)
         trans_table[(int) Py_CHARMASK( ((unsigned char*)vdel.buf)[i] )] = -1;
@@ -1435,7 +1443,8 @@ bytes_translate(PyByteArrayObject *self, PyObject *args)
         PyByteArray_Resize(result, output - output_start);
 
 done:
-    PyBuffer_Release(&vtable);
+    if (tableobj != NULL)
+        PyBuffer_Release(&vtable);
     if (delobj != NULL)
         PyBuffer_Release(&vdel);
     return result;