Added support for per-cursor and per-connection typecasters.

Author: fogzot

ChangeLog

@@ -1,3 +1,7 @@
+2007-02-22  Federico Di Gregorio  <fog@initd.org>
+
+	* Added support for per-connection and per-cursor typecasters.
+
 2007-02-11  Federico Di Gregorio  <fog@initd.org>
 
 	* psycopg/pqpath.c: ported psycopg1 patch from #135.

examples/typecast.py

@@ -0,0 +1,67 @@
+# typecast.py - example of per-cursor and per-connection typecasters.
+#
+# Copyright (C) 2001-2007 Federico Di Gregorio  <fog@debian.org>
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 2, or (at your option) any later
+# version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTIBILITY
+# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+# for more details.
+
+## put in DSN your DSN string
+
+DSN = 'dbname=test'
+
+## don't modify anything below this line (except for experimenting)
+
+class SimpleQuoter(object):
+    def sqlquote(x=None):
+        return "'bar'"
+
+import sys
+import psycopg2
+import psycopg2.extensions
+
+if len(sys.argv) > 1:
+    DSN = sys.argv[1]
+
+print "Opening connection using dns:", DSN
+conn = psycopg2.connect(DSN)
+print "Encoding for this connection is", conn.encoding
+
+curs = conn.cursor()
+curs.execute("SELECT 'text'::text AS foo")
+textoid = curs.description[0][1]
+print "Oid for the text datatype is", textoid
+
+def castA(s, curs):
+    if s is not None: return "(A) " + s
+TYPEA = psycopg2.extensions.new_type((textoid,), "TYPEA", castA)
+
+def castB(s, curs):
+    if s is not None: return "(B) " + s
+TYPEB = psycopg2.extensions.new_type((textoid,), "TYPEB", castB)
+
+curs = conn.cursor()
+curs.execute("SELECT 'some text.'::text AS foo")
+print "Some text from plain connection:", curs.fetchone()[0]
+
+psycopg2.extensions.register_type(TYPEA, conn)
+curs = conn.cursor()
+curs.execute("SELECT 'some text.'::text AS foo")
+print "Some text from connection with typecaster:", curs.fetchone()[0]
+
+curs = conn.cursor()
+psycopg2.extensions.register_type(TYPEB, curs)
+curs.execute("SELECT 'some text.'::text AS foo")
+print "Some text from cursor with typecaster:", curs.fetchone()[0]
+
+curs = conn.cursor()
+curs.execute("SELECT 'some text.'::text AS foo")
+print "Some text from connection with typecaster again:", curs.fetchone()[0]
+
+

psycopg/connection.h

@@ -74,7 +74,11 @@ typedef struct {
     PyObject *exc_IntegrityError;
     PyObject *exc_DataError;
     PyObject *exc_NotSupportedError;
-    
+
+    /* per-connection typecasters */
+    PyObject *string_types;   /* a set of typecasters for string types */
+    PyObject *binary_types;   /* a set of typecasters for binary types */
+
 } connectionObject;
 
 /* C-callable functions in connection_int.c and connection_ext.c */

psycopg/connection_type.c

@@ -276,7 +276,11 @@ static struct PyMemberDef connectionObject_members[] = {
         "The current connection string."},
     {"status", T_LONG,
         offsetof(connectionObject, status), RO,
-	"The current transaction status."},
+	    "The current transaction status."},
+    {"string_types", T_OBJECT, offsetof(connectionObject, string_types), RO,
+        "A set of typecasters to convert textual values."},
+    {"binary_types", T_OBJECT, offsetof(connectionObject, binary_types), RO,
+        "A set of typecasters to convert binary values."},
 #endif    
     {NULL}
 };
@@ -301,7 +305,8 @@ connection_setup(connectionObject *self, char *dsn)
     self->async_cursor = NULL;
     self->pgconn = NULL;
     self->mark = 0;
-   
+    self->string_types = PyDict_New();   
+    self->binary_types = PyDict_New();
 
     pthread_mutex_init(&(self->lock), NULL);
 
@@ -339,6 +344,8 @@ connection_dealloc(PyObject* obj)
     Py_XDECREF(self->notice_list);
     Py_XDECREF(self->notifies);
     Py_XDECREF(self->async_cursor);
+    Py_XDECREF(self->string_types);
+    Py_XDECREF(self->binary_types);
 
     pthread_mutex_destroy(&(self->lock));
 

psycopg/cursor_type.c

@@ -1448,7 +1448,9 @@ cursor_dealloc(PyObject* obj)
     Py_XDECREF(self->tuple_factory);
     Py_XDECREF(self->tzinfo_factory);
     Py_XDECREF(self->query);
-
+    Py_XDECREF(self->string_types);
+    Py_XDECREF(self->binary_types);
+    
     IFCLEARPGRES(self->pgres);
 
     Dprintf("cursor_dealloc: deleted cursor object at %p, refcnt = %d",

psycopg/pqpath.c

@@ -501,34 +501,44 @@ _pq_fetch_tuples(cursorObject *curs)
 
         PyTuple_SET_ITEM(curs->description, i, dtitem);
 
-        /* fill the right cast function by accessing the global dictionary of
-           casting objects.  if we got no defined cast use the default
-           one */
-        if (!(cast = PyDict_GetItem(curs->casts, type))) {
-            Dprintf("_pq_fetch_tuples: cast %d not in per-cursor dict", ftype);
-            if (!(cast = PyDict_GetItem(psyco_types, type))) {
-                Dprintf("_pq_fetch_tuples: cast %d not found, using default",
-                        PQftype(curs->pgres,i));
-                cast = psyco_default_cast;
-            }
+        /* fill the right cast function by accessing three different dictionaries:
+           - the per-cursor dictionary, if available (can be NULL or None)
+           - the per-connection dictionary (always exists but can be null)
+           - the global dictionary (at module level)
+           if we get no defined cast use the default one */
+
+        Dprintf("_pq_fetch_tuples: looking for cast %d:", ftype);
+        if (curs->string_types != NULL && curs->string_types != Py_None) {
+            cast = PyDict_GetItem(curs->string_types, type);
+            Dprintf("_pq_fetch_tuples:     per-cursor dict: %p", cast);
+        }
+        if (cast == NULL) {
+            cast = PyDict_GetItem(curs->conn->string_types, type);
+            Dprintf("_pq_fetch_tuples:     per-connection dict: %p", cast);
+        }
+        if (cast == NULL) {
+            cast = PyDict_GetItem(psyco_types, type);
+            Dprintf("_pq_fetch_tuples:     global dict: %p", cast);
         }
+        if (cast == NULL) cast = psyco_default_cast;
+        
         /* else if we got binary tuples and if we got a field that
            is binary use the default cast
            FIXME: what the hell am I trying to do here? This just can't work..
         */
-        else if (pgbintuples && cast == psyco_default_binary_cast) {
+        if (pgbintuples && cast == psyco_default_binary_cast) {
             Dprintf("_pq_fetch_tuples: Binary cursor and "
                     "binary field: %i using default cast",
                     PQftype(curs->pgres,i));
             cast = psyco_default_cast;
         }
+
         Dprintf("_pq_fetch_tuples: using cast at %p (%s) for type %d",
                 cast, PyString_AS_STRING(((typecastObject*)cast)->name),
                 PQftype(curs->pgres,i));
         Py_INCREF(cast);
         PyTuple_SET_ITEM(curs->casts, i, cast);
 
-
         /* 1/ fill the other fields */
         PyTuple_SET_ITEM(dtitem, 0,
                          PyString_FromString(PQfname(curs->pgres, i)));

psycopg/psycopgmodule.c

@@ -228,17 +228,40 @@ psyco_connect(PyObject *self, PyObject *args, PyObject *keywds)
 "    the string representation returned by PostgreSQL (`None` if ``NULL``)\n" \
 "    and ``cur`` is the cursor from which data are read."
 
+static void
+_psyco_register_type_set(PyObject **dict, PyObject *type)
+{
+    if (*dict == NULL)
+        *dict = PyDict_New();
+    typecast_add(type, *dict, 0);
+}
+
 static PyObject *
 psyco_register_type(PyObject *self, PyObject *args)
 {
-    PyObject *type;
+    PyObject *type, *obj;
 
-    if (!PyArg_ParseTuple(args, "O!", &typecastType, &type)) {
+    if (!PyArg_ParseTuple(args, "O!|O", &typecastType, &type, &obj)) {
         return NULL;
     }
 
-    typecast_add(type, 0);
-    
+    if (obj != NULL) {
+        if (obj->ob_type == &cursorType) {
+            _psyco_register_type_set(&(((cursorObject*)obj)->string_types), type);
+        }
+        else if (obj->ob_type == &connectionType) {
+            typecast_add(type, ((connectionObject*)obj)->string_types, 0);
+        }
+        else {
+            PyErr_SetString(PyExc_TypeError,
+                "argument 2 must be a connection, cursor or None");
+            return NULL;
+        }
+    }
+    else {
+        typecast_add(type, NULL, 0);
+    }
+
     Py_INCREF(Py_None);
     return Py_None;
 }

psycopg/typecast.c

@@ -226,7 +226,7 @@ typecast_init(PyObject *dict)
 
         t = (typecastObject *)typecast_from_c(&(typecast_builtins[i]), dict);
         if (t == NULL) return -1;
-        if (typecast_add((PyObject *)t, 0) != 0) return -1;
+        if (typecast_add((PyObject *)t, NULL, 0) != 0) return -1;
 
         PyDict_SetItem(dict, t->name, (PyObject *)t);
 
@@ -264,7 +264,7 @@ typecast_init(PyObject *dict)
 
 /* typecast_add - add a type object to the dictionary */
 int
-typecast_add(PyObject *obj, int binary)
+typecast_add(PyObject *obj, PyObject *dict, int binary)
 {
     PyObject *val;
     Py_ssize_t len, i;
@@ -274,16 +274,14 @@ typecast_add(PyObject *obj, int binary)
     Dprintf("typecast_add: object at %p, values refcnt = %d",
             obj, type->values->ob_refcnt);
 
+    if (dict == NULL)
+        dict = (binary ? psyco_binary_types : psyco_types);
+
     len = PyTuple_Size(type->values);
     for (i = 0; i < len; i++) {
         val = PyTuple_GetItem(type->values, i);
         Dprintf("typecast_add:     adding val: %ld", PyInt_AsLong(val));
-        if (binary) {
-            PyDict_SetItem(psyco_binary_types, val, obj);
-        }
-        else {
-            PyDict_SetItem(psyco_types, val, obj);
-        }
+        PyDict_SetItem(dict, val, obj);
     }
 
     Dprintf("typecast_add:     base caster: %p", type->bcast);

psycopg/typecast.h

@@ -69,7 +69,7 @@ extern PyObject *psyco_default_binary_cast;
 
 /* used by module.c to init the type system and register types */
 extern int typecast_init(PyObject *dict);
-extern int typecast_add(PyObject *obj, int binary);
+extern int typecast_add(PyObject *obj, PyObject *dict, int binary);
 
 /* the C callable typecastObject creator function */
 extern PyObject *typecast_from_c(typecastObject_initlist *type, PyObject *d);

setup.cfg

@@ -1,5 +1,5 @@
 [build_ext]
-define=PSYCOPG_DEBUG,PSYCOPG_EXTENSIONS,PSYCOPG_DISPLAY_SIZE,PSYCOPG_NEW_BOOLEAN,HAVE_PQFREEMEM,HAVE_PQPROTOCOL3
+define=PSYCOPG_EXTENSIONS,PSYCOPG_DISPLAY_SIZE,PSYCOPG_NEW_BOOLEAN,HAVE_PQFREEMEM,HAVE_PQPROTOCOL3
 # PSYCOPG_EXTENSIONS enables extensions to PEP-249 (you really want this)
 # PSYCOPG_DISPLAY_SIZE enable display size calculation (a little slower)
 # HAVE_PQFREEMEM should be defined on PostgreSQL >= 7.4