psycopg2/psycopg/adapter_binary.c

/* adapter_binary.c - Binary objects
 *
 * Copyright (C) 2003 Federico Di Gregorio <fog@debian.org>
 *
 * This file is part of psycopg.
 *
 * 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
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 */

#include <Python.h>
#include <structmember.h>
#include <stringobject.h>

#include <libpq-fe.h>
#include <string.h>

#define PSYCOPG_MODULE
#include "psycopg/config.h"
#include "psycopg/python.h"
#include "psycopg/psycopg.h"
#include "psycopg/adapter_binary.h"
#include "psycopg/microprotocols_proto.h"

/** the quoting code */

#ifndef PSYCOPG_OWN_QUOTING
#define binary_escape PQescapeBytea
#else
static unsigned char *
binary_escape(char *from, int from_length, int *to_length)
{
    unsigneed char *quoted, *chptr, *newptr;
    int i, space, new_space;

    space = from_length + 2;

    Py_BEGIN_ALLOW_THREADS;

    quoted = (unsigned char*)calloc(space, sizeof(char));
    if (quoted == NULL) return NULL;

    chptr = quoted;

    for (i=0; i < len; i++) {
        if (chptr - quoted > space - 6) {
            new_space  =  space * ((space) / (i + 1)) + 2 + 6;
            if (new_space - space < 1024) space += 1024;
            else space = new_space;
            newptr = (unsigned char *)realloc(quoted, space);
            if (newptr == NULL) {
                free(quoted);
                return NULL;
            }
            /* chptr has to be moved to the new location*/
            chptr = newptr + (chptr - quoted);
            quoted = newptr;
            Dprintf("binary_escape: reallocated %i bytes at %p", space,quoted);
        }
        if (from[i]) {
            if (from[i] >= ' ' && from[i] <= '~') {
                if (from[i] == '\'') {
                    *chptr = '\\';
                    chptr++;
                    *chptr = '\'';
                    chptr++;
                }
                else if (from[i] == '\\') {
                    memcpy(chptr, "\\\\\\\\", 4);
                    chptr += 4;
                }
                else {
                    /* leave it as it is if ascii printable */
                    *chptr = from[i];
                    chptr++;
                }
            }
            else {
                unsigned char c;
                
                /* escape to octal notation \nnn */
                *chptr++ = '\\';
                *chptr++ = '\\';
                c = from[i];
                *chptr = ((c >> 6) & 0x07) + 0x30; chptr++;
                *chptr = ((c >> 3) & 0x07) + 0x30; chptr++;
                *chptr = ( c       & 0x07) + 0x30; chptr++;
            }
        }
        else {
            /* escape null as \\000 */
            memcpy(chptr, "\\\\000", 5);
            chptr += 5;
        }
    }
    *chptr = '\0';

    Py_END_ALLOW_THREADS;

    *to_size = chptr - quoted + 1;
    return quoted;
}
#endif

/* binary_quote - do the quote process on plain and unicode strings */

static PyObject *
binary_quote(binaryObject *self)
{
    char *to;
    const char *buffer;
    int buffer_len;
    size_t len = 0;
    
    /* if we got a plain string or a buffer we escape it and save the buffer */
    if (PyString_Check(self->wrapped) || PyBuffer_Check(self->wrapped)) {
        /* escape and build quoted buffer */
        PyObject_AsCharBuffer(self->wrapped, &buffer, &buffer_len);
        to = (char *)binary_escape(buffer, buffer_len, &len);
        if (to == NULL) {
            PyErr_NoMemory();
            return NULL;
        }

        self->buffer = PyString_FromFormat("'%s'", to);
        PQfreemem(to);
    }
    
    /* if the wrapped object is not a string or a buffer, this is an error */ 
    else {
        PyErr_SetString(PyExc_TypeError, "can't escape non-string object");
        return NULL;
    }
    
    return self->buffer;
}

/* binary_str, binary_getquoted - return result of quoting */

static PyObject *
binary_str(binaryObject *self)
{
    if (self->buffer == NULL) {
        binary_quote(self);
    }
    Py_XINCREF(self->buffer);
    return self->buffer;
}

PyObject *
binary_getquoted(binaryObject *self, PyObject *args)
{
    if (!PyArg_ParseTuple(args, "")) return NULL;
    return binary_str(self);
}

PyObject *
binary_conform(binaryObject *self, PyObject *args)
{
    PyObject *res, *proto;
    
    if (!PyArg_ParseTuple(args, "O", &proto)) return NULL;

    if (proto == (PyObject*)&isqlquoteType)
        res = (PyObject*)self;
    else
        res = Py_None;
    
    Py_INCREF(res);
    return res;
}

/** the Binary object **/

/* object member list */

static struct PyMemberDef binaryObject_members[] = {
    {"adapted", T_OBJECT, offsetof(binaryObject, wrapped), RO},
    {"buffer", T_OBJECT, offsetof(binaryObject, buffer), RO},
    {NULL}
};

/* object method table */

static PyMethodDef binaryObject_methods[] = {
    {"getquoted", (PyCFunction)binary_getquoted, METH_VARARGS,
     "getquoted() -> wrapped object value as SQL-quoted binary string"},
    {"__conform__", (PyCFunction)binary_conform, METH_VARARGS, NULL},
    {NULL}  /* Sentinel */
};

/* initialization and finalization methods */

static int
binary_setup(binaryObject *self, PyObject *str)
{
    Dprintf("binary_setup: init binary object at %p, refcnt = %d",
            self, ((PyObject *)self)->ob_refcnt);

    self->buffer = NULL;
    self->wrapped = str;
    Py_INCREF(self->wrapped);
    
    Dprintf("binary_setup: good binary object at %p, refcnt = %d",
            self, ((PyObject *)self)->ob_refcnt);
    return 0;
}

static void
binary_dealloc(PyObject* obj)
{
    binaryObject *self = (binaryObject *)obj;

    Py_XDECREF(self->wrapped);
    Py_XDECREF(self->buffer);
    
    Dprintf("binary_dealloc: deleted binary object at %p, refcnt = %d",
            obj, obj->ob_refcnt);
    
    obj->ob_type->tp_free(obj);
}

static int
binary_init(PyObject *obj, PyObject *args, PyObject *kwds)
{
    PyObject *str;
    
    if (!PyArg_ParseTuple(args, "O", &str))
        return -1;

    return binary_setup((binaryObject *)obj, str);
}

static PyObject *
binary_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{    
    return type->tp_alloc(type, 0);
}

static void
binary_del(PyObject* self)
{
    PyObject_Del(self);
}

static PyObject *
binary_repr(binaryObject *self)
{
    return PyString_FromFormat("<psycopg.Binary object at %p>", self);
}

/* object type */

#define binaryType_doc \
"psycopg.Binary(buffer) -> new binary object"

PyTypeObject binaryType = {
    PyObject_HEAD_INIT(NULL)
    0,
    "psycopg._psycopg.Binary",
    sizeof(binaryObject),
    0,
    binary_dealloc, /*tp_dealloc*/
    0,          /*tp_print*/
    0,          /*tp_getattr*/
    0,          /*tp_setattr*/   

    0,          /*tp_compare*/
    (reprfunc)binary_repr, /*tp_repr*/
    0,          /*tp_as_number*/
    0,          /*tp_as_sequence*/
    0,          /*tp_as_mapping*/
    0,          /*tp_hash */

    0,          /*tp_call*/
    (reprfunc)binary_str, /*tp_str*/
    0,          /*tp_getattro*/
    0,          /*tp_setattro*/
    0,          /*tp_as_buffer*/

    Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/

    binaryType_doc, /*tp_doc*/
    
    0,          /*tp_traverse*/
    0,          /*tp_clear*/

    0,          /*tp_richcompare*/
    0,          /*tp_weaklistoffset*/

    0,          /*tp_iter*/
    0,          /*tp_iternext*/

    /* Attribute descriptor and subclassing stuff */

    binaryObject_methods, /*tp_methods*/
    binaryObject_members, /*tp_members*/
    0,          /*tp_getset*/
    0,          /*tp_base*/
    0,          /*tp_dict*/
    
    0,          /*tp_descr_get*/
    0,          /*tp_descr_set*/
    0,          /*tp_dictoffset*/
    
    binary_init, /*tp_init*/
    0, /*tp_alloc  will be set to PyType_GenericAlloc in module init*/
    binary_new, /*tp_new*/
    (freefunc)binary_del, /*tp_free  Low-level free-memory routine */
    0,          /*tp_is_gc For PyObject_IS_GC */
    0,          /*tp_bases*/
    0,          /*tp_mro method resolution order */
    0,          /*tp_cache*/
    0,          /*tp_subclasses*/
    0           /*tp_weaklist*/
};


/** module-level functions **/

PyObject *
psyco_Binary(PyObject *module, PyObject *args)
{
    PyObject *str;
    
    if (!PyArg_ParseTuple(args, "O", &str))
        return NULL;
  
    return PyObject_CallFunction((PyObject *)&binaryType, "O", str);
}