Adds support for DB.pget and DBCursor.pget methods.

Based on a patch supplied by Ian Ward <ian@arevco.ca> on the pybsddb
mailing list 2004-03-26.

3 files changed, 225 insertions, 3 deletions

diff --git a/Lib/bsddb/dbobj.py b/Lib/bsddb/dbobj.py
index abda657..3bafafa 100644
--- a/Lib/bsddb/dbobj.py
+++ b/Lib/bsddb/dbobj.py
@@ -134,6 +134,8 @@ class DB(DictMixin):
         return apply(self._cobj.fd, args, kwargs)
     def get(self, *args, **kwargs):
         return apply(self._cobj.get, args, kwargs)
+    def pget(self, *args, **kwargs):
+        return apply(self._cobj.pget, args, kwargs)
     def get_both(self, *args, **kwargs):
         return apply(self._cobj.get_both, args, kwargs)
     def get_byteswapped(self, *args, **kwargs):
diff --git a/Lib/bsddb/test/test_associate.py b/Lib/bsddb/test/test_associate.py
index fc92c22..7be5ba0 100644
--- a/Lib/bsddb/test/test_associate.py
+++ b/Lib/bsddb/test/test_associate.py
@@ -83,6 +83,7 @@ musicdata = {
 52: ("David Lanz", "Cristofori's Dream", "New Age"),
 53: ("David Lanz", "Heartsounds", "New Age"),
 54: ("David Lanz", "Leaves on the Seine", "New Age"),
+99: ("unknown artist", "Unnamed song", "Unknown"),
 }
 
 #----------------------------------------------------------------------
@@ -117,6 +118,7 @@ class AssociateTestCase(unittest.TestCase):
 
     def createDB(self):
         self.primary = db.DB(self.env)
+        self.primary.set_get_returns_none(2)
         self.primary.open(self.filename, "primary", self.dbtype,
                           db.DB_CREATE | db.DB_THREAD)
 
@@ -136,6 +138,7 @@ class AssociateTestCase(unittest.TestCase):
 
         secDB = db.DB(self.env)
         secDB.set_flags(db.DB_DUP)
+        secDB.set_get_returns_none(2)
         secDB.open(self.filename, "secondary", db.DB_BTREE,
                    db.DB_CREATE | db.DB_THREAD)
         self.getDB().associate(secDB, self.getGenre)
@@ -166,6 +169,16 @@ class AssociateTestCase(unittest.TestCase):
 
 
     def finish_test(self, secDB):
+        # 'Blues' should not be in the secondary database
+        vals = secDB.pget('Blues')
+        assert vals == None, vals
+
+        vals = secDB.pget('Unknown')
+        assert vals[0] == 99 or vals[0] == '99', vals
+        vals[1].index('Unknown')
+        vals[1].index('Unnamed')
+        vals[1].index('unknown')
+
         if verbose:
             print "Primary key traversal:"
         c = self.getDB().cursor()
@@ -187,6 +200,18 @@ class AssociateTestCase(unittest.TestCase):
             print "Secondary key traversal:"
         c = secDB.cursor()
         count = 0
+
+        # test cursor pget
+        vals = c.pget('Unknown', flags=db.DB_LAST)
+        assert vals[1] == 99 or vals[1] == '99', vals
+        assert vals[0] == 'Unknown'
+        vals[2].index('Unknown')
+        vals[2].index('Unnamed')
+        vals[2].index('unknown')
+
+        vals = c.pget('Unknown', data='wrong value', flags=db.DB_GET_BOTH)
+        assert vals == None, vals
+
         rec = c.first()
         assert rec[0] == "Jazz"
         while rec is not None:
diff --git a/Modules/_bsddb.c b/Modules/_bsddb.c
index e00a64c..b7f6518 100644
--- a/Modules/_bsddb.c
+++ b/Modules/_bsddb.c
@@ -97,7 +97,7 @@
 #error "eek! DBVER can't handle minor versions > 9"
 #endif
 
-#define PY_BSDDB_VERSION "4.2.7"
+#define PY_BSDDB_VERSION "4.2.8"
 static char *rcs_id = "$Id$";
 
 
@@ -1463,6 +1463,94 @@ DB_get(DBObject* self, PyObject* args, PyObject* kwargs)
     return retval;
 }
 
+static PyObject*
+DB_pget(DBObject* self, PyObject* args, PyObject* kwargs)
+{
+    int err, flags=0;
+    PyObject* txnobj = NULL;
+    PyObject* keyobj;
+    PyObject* dfltobj = NULL;
+    PyObject* retval = NULL;
+    int dlen = -1;
+    int doff = -1;
+    DBT key, pkey, data;
+    DB_TXN *txn = NULL;
+    char* kwnames[] = {"key", "default", "txn", "flags", "dlen", "doff", NULL};
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|OOiii:pget", kwnames,
+                                     &keyobj, &dfltobj, &txnobj, &flags, &dlen,
+                                     &doff))
+        return NULL;
+
+    CHECK_DB_NOT_CLOSED(self);
+    if (!make_key_dbt(self, keyobj, &key, &flags))
+        return NULL;
+    if (!checkTxnObj(txnobj, &txn)) {
+        FREE_DBT(key);
+        return NULL;
+    }
+
+    CLEAR_DBT(data);
+    if (CHECK_DBFLAG(self, DB_THREAD)) {
+        /* Tell BerkeleyDB to malloc the return value (thread safe) */
+        data.flags = DB_DBT_MALLOC;
+    }
+    if (!add_partial_dbt(&data, dlen, doff)) {
+        FREE_DBT(key);
+        return NULL;
+    }
+
+    CLEAR_DBT(pkey);
+    pkey.flags = DB_DBT_MALLOC;
+    
+    MYDB_BEGIN_ALLOW_THREADS;
+    err = self->db->pget(self->db, txn, &key, &pkey, &data, flags);
+    MYDB_END_ALLOW_THREADS;
+
+    if ((err == DB_NOTFOUND) && (dfltobj != NULL)) {
+        err = 0;
+        Py_INCREF(dfltobj);
+        retval = dfltobj;
+    }
+    else if ((err == DB_NOTFOUND) && self->moduleFlags.getReturnsNone) {
+        err = 0;
+        Py_INCREF(Py_None);
+        retval = Py_None;
+    }
+    else if (!err) {
+        PyObject *pkeyObj;
+        PyObject *dataObj;
+        dataObj = PyString_FromStringAndSize(data.data, data.size);
+
+        if (self->primaryDBType == DB_RECNO ||
+            self->primaryDBType == DB_QUEUE)
+            pkeyObj = PyInt_FromLong(*(long *)pkey.data);
+        else
+            pkeyObj = PyString_FromStringAndSize(pkey.data, pkey.size);
+
+        if (flags & DB_SET_RECNO) /* return key , pkey and data */
+        {
+            PyObject *keyObj;
+            int type = _DB_get_type(self);
+            if (type == DB_RECNO || type == DB_QUEUE)
+                keyObj = PyInt_FromLong(*(long *)key.data);
+            else
+                keyObj = PyString_FromStringAndSize(key.data, key.size);
+            retval = Py_BuildValue("OOO", keyObj, pkeyObj, dataObj);
+        }
+        else /* return just the pkey and data */
+        {
+            retval = Py_BuildValue("OO", pkeyObj, dataObj);
+        }
+	FREE_DBT(pkey);
+        FREE_DBT(data);
+    }
+    FREE_DBT(key);
+
+    RETURN_IF_ERR();
+    return retval;
+}
+
 
 /* Return size of entry */
 static PyObject*
@@ -2827,6 +2915,106 @@ DBC_get(DBCursorObject* self, PyObject* args, PyObject *kwargs)
     return retval;
 }
 
+static PyObject*
+DBC_pget(DBCursorObject* self, PyObject* args, PyObject *kwargs)
+{
+    int err, flags=0;
+    PyObject* keyobj = NULL;
+    PyObject* dataobj = NULL;
+    PyObject* retval = NULL;
+    int dlen = -1;
+    int doff = -1;
+    DBT key, pkey, data;
+    char* kwnames[] = { "key","data", "flags", "dlen", "doff", NULL };
+
+    CLEAR_DBT(key);
+    CLEAR_DBT(data);
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "i|ii:pget", &kwnames[2],
+				     &flags, &dlen, &doff))
+    {
+        PyErr_Clear();
+        if (!PyArg_ParseTupleAndKeywords(args, kwargs, "Oi|ii:pget",
+                                         &kwnames[1], 
+					 &keyobj, &flags, &dlen, &doff))
+        {
+            PyErr_Clear();
+            if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OOi|ii:pget",
+                                             kwnames, &keyobj, &dataobj,
+                                             &flags, &dlen, &doff))
+            {
+                return NULL;
+	    }
+	}
+    }
+
+    CHECK_CURSOR_NOT_CLOSED(self);
+
+    if (keyobj && !make_key_dbt(self->mydb, keyobj, &key, NULL))
+        return NULL;
+    if ( (dataobj && !make_dbt(dataobj, &data)) ||
+         (!add_partial_dbt(&data, dlen, doff)) ) {
+        FREE_DBT(key);
+        return NULL;
+    }
+
+    if (CHECK_DBFLAG(self->mydb, DB_THREAD)) {
+        data.flags = DB_DBT_MALLOC;
+        if (!(key.flags & DB_DBT_REALLOC)) {
+            key.flags |= DB_DBT_MALLOC;
+        }
+    }
+
+    CLEAR_DBT(pkey);
+    pkey.flags = DB_DBT_MALLOC;
+
+    MYDB_BEGIN_ALLOW_THREADS;
+    err = self->dbc->c_pget(self->dbc, &key, &pkey, &data, flags);
+    MYDB_END_ALLOW_THREADS;
+
+    if ((err == DB_NOTFOUND) && self->mydb->moduleFlags.getReturnsNone) {
+        Py_INCREF(Py_None);
+        retval = Py_None;
+    }
+    else if (makeDBError(err)) {
+        retval = NULL;
+    }
+    else {
+        PyObject *pkeyObj;
+        PyObject *dataObj;
+        dataObj = PyString_FromStringAndSize(data.data, data.size);
+
+        if (self->mydb->primaryDBType == DB_RECNO ||
+            self->mydb->primaryDBType == DB_QUEUE)
+            pkeyObj = PyInt_FromLong(*(long *)pkey.data);
+        else
+            pkeyObj = PyString_FromStringAndSize(pkey.data, pkey.size);
+
+        if (flags & DB_SET_RECNO) /* return key, pkey and data */
+        {
+            PyObject *keyObj;
+            int type = _DB_get_type(self->mydb);
+            if (type == DB_RECNO || type == DB_QUEUE)
+                keyObj = PyInt_FromLong(*(long *)key.data);
+            else
+                keyObj = PyString_FromStringAndSize(key.data, key.size);
+            retval = Py_BuildValue("OOO", keyObj, pkeyObj, dataObj);
+            FREE_DBT(key);
+        }
+        else /* return just the pkey and data */
+        {
+            retval = Py_BuildValue("OO", pkeyObj, dataObj);
+        }
+        FREE_DBT(pkey);
+        FREE_DBT(data);
+    }
+    /* the only time REALLOC should be set is if we used an integer
+     * key that make_key_dbt malloc'd for us.  always free these. */
+    if (key.flags & DB_DBT_REALLOC) {
+        FREE_DBT(key);
+    }
+    return retval;
+}
+
 
 static PyObject*
 DBC_get_recno(DBCursorObject* self, PyObject* args)
@@ -2974,8 +3162,13 @@ DBC_set(DBCursorObject* self, PyObject* args, PyObject *kwargs)
             break;
         }
         FREE_DBT(data);
+        FREE_DBT(key);
+    }
+    /* the only time REALLOC should be set is if we used an integer
+     * key that make_key_dbt malloc'd for us.  always free these. */
+    if (key.flags & DB_DBT_REALLOC) {
+        FREE_DBT(key);
     }
-    FREE_DBT(key);
 
     return retval;
 }
@@ -3044,7 +3237,7 @@ DBC_set_range(DBCursorObject* self, PyObject* args, PyObject* kwargs)
         FREE_DBT(data);
     }
     /* the only time REALLOC should be set is if we used an integer
-     * key that make_dbt_key malloc'd for us.  always free these. */
+     * key that make_key_dbt malloc'd for us.  always free these. */
     if (key.flags & DB_DBT_REALLOC) {
         FREE_DBT(key);
     }
@@ -4183,6 +4376,7 @@ static PyMethodDef DB_methods[] = {
     {"delete",          (PyCFunction)DB_delete,         METH_VARARGS|METH_KEYWORDS},
     {"fd",              (PyCFunction)DB_fd,             METH_VARARGS},
     {"get",             (PyCFunction)DB_get,            METH_VARARGS|METH_KEYWORDS},
+    {"pget",            (PyCFunction)DB_pget,           METH_VARARGS|METH_KEYWORDS},
     {"get_both",        (PyCFunction)DB_get_both,       METH_VARARGS|METH_KEYWORDS},
     {"get_byteswapped", (PyCFunction)DB_get_byteswapped,METH_VARARGS},
     {"get_size",        (PyCFunction)DB_get_size,       METH_VARARGS|METH_KEYWORDS},
@@ -4242,6 +4436,7 @@ static PyMethodDef DBCursor_methods[] = {
     {"dup",             (PyCFunction)DBC_dup,           METH_VARARGS},
     {"first",           (PyCFunction)DBC_first,         METH_VARARGS|METH_KEYWORDS},
     {"get",             (PyCFunction)DBC_get,           METH_VARARGS|METH_KEYWORDS},
+    {"pget",            (PyCFunction)DBC_pget,          METH_VARARGS|METH_KEYWORDS},
     {"get_recno",       (PyCFunction)DBC_get_recno,     METH_VARARGS},
     {"last",            (PyCFunction)DBC_last,          METH_VARARGS|METH_KEYWORDS},
     {"next",            (PyCFunction)DBC_next,          METH_VARARGS|METH_KEYWORDS},