Array quoting (not completely working.)

Author: fogzot

psycopg/typecast_array.c

@@ -22,20 +22,27 @@
 
 /** typecast_array_scan - scan a string looking for array items **/
 
-#define ASCAN_EOF   0
-#define ASCAN_BEGIN 1
-#define ASCAN_END   2
-#define ASCAN_TOKEN 3
+#define ASCAN_EOF    0
+#define ASCAN_BEGIN  1
+#define ASCAN_END    2
+#define ASCAN_TOKEN  3
+#define ASCAN_QUOTED 4
 
 static int
 typecast_array_tokenize(unsigned char *str, int strlength,
                         int *pos, unsigned char** token, int *length)
 {
-    int i = *pos;
+    int i;
+    int quoted = 0;
+
+    /* first we check for quotes, used when the content of the item contains
+       special or quoted characters */
+    if (str[*pos] == '"') {
+        quoted = 1;
+        *pos += 1;
+    }
 
-    Dprintf("TOKENIZE for %d, pos = %d", strlength, *pos);
-    
-    while (i < strlength) {
+    for (i = *pos ; i < strlength ; i++) {
         switch (str[i]) {
         case '{':
             *pos = i+1;
@@ -48,17 +55,21 @@ typecast_array_tokenize(unsigned char *str, int strlength,
         case ',':
             *token = &str[*pos];
             *length = i - *pos;
+            if (quoted == 1)
+                *length -= 1;
             *pos = i+1;
-            Dprintf("TOKENIZE pos = %d, length = %d", *pos, *length);
             return ASCAN_TOKEN;
 
         default:
-            i++;
+            /* nothing to do right now */
+            break;
         }
     }
 
     *token = &str[*pos];
     *length = i - *pos;
+    if (quoted == 1)
+        *length -= 1;
 
     return ASCAN_EOF;
 }
@@ -89,10 +100,12 @@ typecast_array_scan(unsigned char *str, int strlength,
     return 1;
 }
 
-/** LONGINTEGERARRAY and INTEGERARRAY - cast integers arrays **/
 
+/** GENERIC - a generic typecaster that can be used when no special actions
+    have to be taken on the single items **/
+   
 static PyObject *
-typecast_INTEGERARRAY_cast(unsigned char *str, int len, PyObject *curs)
+typecast_GENERIC_ARRAY_cast(unsigned char *str, int len, PyObject *curs)
 {
     PyObject *obj = NULL;
     PyObject *base = ((typecastObject*)((cursorObject*)curs)->caster)->bcast;
@@ -114,16 +127,12 @@ typecast_INTEGERARRAY_cast(unsigned char *str, int len, PyObject *curs)
     return obj;
 }
 
-#define typecast_LONGINTEGERARRAY_cast typecast_INTEGERARRAY_cast
+/** LONGINTEGERARRAY and INTEGERARRAY - cast integers arrays **/
+
+#define typecast_LONGINTEGERARRAY_cast typecast_GENERIC_ARRAY_cast
+#define typecast_INTEGERARRAY_cast typecast_GENERIC_ARRAY_cast
 
 /** STRINGARRAY - cast integers arrays **/
 
-static PyObject *
-typecast_STRINGARRAY_cast(unsigned char *str, int len, PyObject *curs)
-{
-    PyObject* obj = NULL;
-     
-    if (str == NULL) {Py_INCREF(Py_None); return Py_None;}
+#define typecast_STRINGARRAY_cast typecast_GENERIC_ARRAY_cast
 
-    return obj;
-}

psycopg/typecast_basic.c

@@ -30,8 +30,11 @@ typecast_INTEGER_cast(unsigned char *s, int len, PyObject *curs)
     unsigned char buffer[12];
 
     if (s == NULL) {Py_INCREF(Py_None); return Py_None;}
-    strncpy(buffer, s, len); buffer[len] = '\0';
-    return PyInt_FromString(buffer, NULL, 0);
+    if (s[len] != '\0') {
+        strncpy(buffer, s, len); buffer[len] = '\0';
+        s = buffer;
+    }
+    return PyInt_FromString(s, NULL, 0);
 }
 
 /** LONGINTEGER - cast long integers (8 bytes) to python long **/
@@ -42,7 +45,10 @@ typecast_LONGINTEGER_cast(unsigned char *s, int len, PyObject *curs)
     unsigned char buffer[24];
 
     if (s == NULL) {Py_INCREF(Py_None); return Py_None;}
-    strncpy(buffer, s, len); buffer[len] = '\0';
+    if (s[len] != '\0') {
+        strncpy(buffer, s, len); buffer[len] = '\0';
+        s = buffer;
+    }
     return PyLong_FromString(s, NULL, 0);
 }
 
@@ -51,10 +57,14 @@ typecast_LONGINTEGER_cast(unsigned char *s, int len, PyObject *curs)
 static PyObject *
 typecast_FLOAT_cast(unsigned char *s, int len, PyObject *curs)
 {
+    /* FIXME: is 64 large enough for any float? */
     unsigned char buffer[64];
 
     if (s == NULL) {Py_INCREF(Py_None); return Py_None;}
-    strncpy(buffer, s, len); buffer[len] = '\0';
+    if (s[len] != '\0') {
+        strncpy(buffer, s, len); buffer[len] = '\0';
+        s = buffer;
+    }
     return PyFloat_FromDouble(atof(s));
 }
 
@@ -144,17 +154,23 @@ static PyObject *
 typecast_BINARY_cast(unsigned char *s, int l, PyObject *curs)
 {
     PyObject *res;
-    unsigned char *str, saved;
+    unsigned char *str, *buffer = NULL;
     size_t len;
 
     if (s == NULL) {Py_INCREF(Py_None); return Py_None;}
 
     /* PQunescapeBytea absolutely wants a 0-terminated string and we don't
-       want to copy the whole buffer, right? Wrong... :/ */
-    saved = s[l]; s[l] = '\0';
+       want to copy the whole buffer, right? Wrong, but there isn't any other
+       way :/ */
+    if (s[l] != '\0') {
+        if ((buffer = PyMem_Malloc(l+1)) == NULL)
+            PyErr_NoMemory();
+        strncpy(buffer, s, l); buffer[l] = '\0';
+        s = buffer;
+    }
     str = PQunescapeBytea(s, &len);
     Dprintf("typecast_BINARY_cast: unescaped %d bytes", len);
-    s[l] = saved;
+    if (buffer) PyMem_Free(buffer);
 
     /* TODO: using a PyBuffer would make this a zero-copy operation but we'll
        need to define our own buffer-derived object to keep a reference to the
@@ -196,7 +212,8 @@ typecast_DECIMAL_cast(unsigned char *s, int len, PyObject *curs)
 
     if (s == NULL) {Py_INCREF(Py_None); return Py_None;}
 
-    buffer = PyMem_Malloc(len+1);
+    if ((buffer = PyMem_Malloc(len+1)) == NULL)
+        PyErr_NoMemory();
     strncpy(buffer, s, len); buffer[len] = '\0';
     res = PyObject_CallFunction(decimalType, "s", buffer);
     PyMem_Free(buffer);

sandbox/array.py

@@ -0,0 +1,15 @@
+import psycopg
+
+conn = psycopg.connect("dbname=test")
+curs = conn.cursor()
+
+curs.execute("SELECT ARRAY[1,2,3] AS foo")
+print curs.fetchone()
+
+curs.execute("SELECT ARRAY['1','2','3'] AS foo")
+print curs.fetchone()
+
+curs.execute("""SELECT ARRAY['','"',''] AS foo""")
+d = curs.fetchone()
+print d, '->', d[0][0], d[0][1], d[0][2]
+