dbmmodule.c 8.46 KB
Newer Older
1 2 3 4

/* DBM module using dictionary interface */


5
#include "Python.h"
6 7 8 9

#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
10 11 12 13 14

/* Some Linux systems install gdbm/ndbm.h, but not ndbm.h.  This supports
 * whichever configure was able to locate.
 */
#if defined(HAVE_NDBM_H)
15
#include <ndbm.h>
16
static char *which_dbm = "ndbm";
17 18
#elif defined(HAVE_DB1_NDBM_H)
#include <db1/ndbm.h>
19
static char *which_dbm = "BSD db";
20 21
#elif defined(HAVE_GDBM_NDBM_H)
#include <gdbm/ndbm.h>
22
static char *which_dbm = "GNU gdbm";
23 24 25
#else
#error "No ndbm.h available!"
#endif
26 27

typedef struct {
28
	PyObject_HEAD
29 30 31 32
	int di_size;	/* -1 means recompute */
	DBM *di_dbm;
} dbmobject;

33
staticforward PyTypeObject Dbmtype;
34 35

#define is_dbmobject(v) ((v)->ob_type == &Dbmtype)
36 37 38
#define check_dbmobject_open(v) if ((v)->di_dbm == NULL) \
               { PyErr_SetString(DbmError, "DBM object has already been closed"); \
                 return NULL; }
39

40
static PyObject *DbmError;
41

42
static PyObject *
43
newdbmobject(char *file, int flags, int mode)
44 45 46
{
        dbmobject *dp;

47
	dp = PyObject_New(dbmobject, &Dbmtype);
48 49 50 51
	if (dp == NULL)
		return NULL;
	dp->di_size = -1;
	if ( (dp->di_dbm = dbm_open(file, flags, mode)) == 0 ) {
52 53 54
		PyErr_SetFromErrno(DbmError);
		Py_DECREF(dp);
		return NULL;
55
	}
56
	return (PyObject *)dp;
57 58 59 60 61
}

/* Methods */

static void
62
dbm_dealloc(register dbmobject *dp)
63 64
{
        if ( dp->di_dbm )
65
		dbm_close(dp->di_dbm);
66
	PyObject_Del(dp);
67 68 69
}

static int
70
dbm_length(dbmobject *dp)
71
{
72 73 74 75
        if (dp->di_dbm == NULL) {
                 PyErr_SetString(DbmError, "DBM object has already been closed"); 
                 return -1; 
        }
76
        if ( dp->di_size < 0 ) {
77 78 79 80 81 82 83 84
		datum key;
		int size;

		size = 0;
		for ( key=dbm_firstkey(dp->di_dbm); key.dptr;
		      key = dbm_nextkey(dp->di_dbm))
			size++;
		dp->di_size = size;
85 86 87 88
	}
	return dp->di_size;
}

89
static PyObject *
90
dbm_subscript(dbmobject *dp, register PyObject *key)
91 92 93
{
	datum drec, krec;
	
94
	if (!PyArg_Parse(key, "s#", &krec.dptr, &krec.dsize) )
95 96
		return NULL;
	
97
        check_dbmobject_open(dp);
98 99
	drec = dbm_fetch(dp->di_dbm, krec);
	if ( drec.dptr == 0 ) {
100 101 102
		PyErr_SetString(PyExc_KeyError,
				PyString_AS_STRING((PyStringObject *)key));
		return NULL;
103 104
	}
	if ( dbm_error(dp->di_dbm) ) {
105 106 107
		dbm_clearerr(dp->di_dbm);
		PyErr_SetString(DbmError, "");
		return NULL;
108
	}
109
	return PyString_FromStringAndSize(drec.dptr, drec.dsize);
110 111 112
}

static int
113
dbm_ass_sub(dbmobject *dp, PyObject *v, PyObject *w)
114 115 116
{
        datum krec, drec;
	
117
        if ( !PyArg_Parse(v, "s#", &krec.dptr, &krec.dsize) ) {
118 119 120
		PyErr_SetString(PyExc_TypeError,
				"dbm mappings have string indices only");
		return -1;
121
	}
122 123 124 125
        if (dp->di_dbm == NULL) {
                 PyErr_SetString(DbmError, "DBM object has already been closed"); 
                 return -1;
        }
126 127
	dp->di_size = -1;
	if (w == NULL) {
128 129 130
		if ( dbm_delete(dp->di_dbm, krec) < 0 ) {
			dbm_clearerr(dp->di_dbm);
			PyErr_SetString(PyExc_KeyError,
131
				      PyString_AS_STRING((PyStringObject *)v));
132 133
			return -1;
		}
134
	} else {
135 136
		if ( !PyArg_Parse(w, "s#", &drec.dptr, &drec.dsize) ) {
			PyErr_SetString(PyExc_TypeError,
137
				     "dbm mappings have string elements only");
138 139 140 141
			return -1;
		}
		if ( dbm_store(dp->di_dbm, krec, drec, DBM_REPLACE) < 0 ) {
			dbm_clearerr(dp->di_dbm);
142
			PyErr_SetString(DbmError,
Fred Drake's avatar
Fred Drake committed
143
					"cannot add item to database");
144 145
			return -1;
		}
146 147
	}
	if ( dbm_error(dp->di_dbm) ) {
148 149 150
		dbm_clearerr(dp->di_dbm);
		PyErr_SetString(DbmError, "");
		return -1;
151 152 153 154
	}
	return 0;
}

155
static PyMappingMethods dbm_as_mapping = {
156 157 158
	(inquiry)dbm_length,		/*mp_length*/
	(binaryfunc)dbm_subscript,	/*mp_subscript*/
	(objobjargproc)dbm_ass_sub,	/*mp_ass_subscript*/
159 160
};

161
static PyObject *
162
dbm__close(register dbmobject *dp, PyObject *args)
163
{
164
	if (!PyArg_ParseTuple(args, ":close"))
165
		return NULL;
166
        if (dp->di_dbm)
167
		dbm_close(dp->di_dbm);
168
	dp->di_dbm = NULL;
169 170
	Py_INCREF(Py_None);
	return Py_None;
171 172
}

173
static PyObject *
174
dbm_keys(register dbmobject *dp, PyObject *args)
175
{
176
	register PyObject *v, *item;
177
	datum key;
178
	int err;
179

180
	if (!PyArg_ParseTuple(args, ":keys"))
181
		return NULL;
182
        check_dbmobject_open(dp);
183
	v = PyList_New(0);
184 185 186
	if (v == NULL)
		return NULL;
	for (key = dbm_firstkey(dp->di_dbm); key.dptr;
187
	     key = dbm_nextkey(dp->di_dbm)) {
188
		item = PyString_FromStringAndSize(key.dptr, key.dsize);
189
		if (item == NULL) {
190
			Py_DECREF(v);
191 192
			return NULL;
		}
193 194
		err = PyList_Append(v, item);
		Py_DECREF(item);
195
		if (err != 0) {
196
			Py_DECREF(v);
197 198
			return NULL;
		}
199 200 201 202
	}
	return v;
}

203
static PyObject *
204
dbm_has_key(register dbmobject *dp, PyObject *args)
205 206 207
{
	datum key, val;
	
208
	if (!PyArg_ParseTuple(args, "s#:has_key", &key.dptr, &key.dsize))
209
		return NULL;
210
        check_dbmobject_open(dp);
211
	val = dbm_fetch(dp->di_dbm, key);
212
	return PyInt_FromLong(val.dptr != NULL);
213 214
}

215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257
static PyObject *
dbm_get(register dbmobject *dp, PyObject *args)
{
	datum key, val;
	PyObject *defvalue = Py_None;

	if (!PyArg_ParseTuple(args, "s#|O:get",
                              &key.dptr, &key.dsize, &defvalue))
		return NULL;
        check_dbmobject_open(dp);
	val = dbm_fetch(dp->di_dbm, key);
	if (val.dptr != NULL)
		return PyString_FromStringAndSize(val.dptr, val.dsize);
	else {
		Py_INCREF(defvalue);
		return defvalue;
	}
}

static PyObject *
dbm_setdefault(register dbmobject *dp, PyObject *args)
{
	datum key, val;
	PyObject *defvalue = NULL;

	if (!PyArg_ParseTuple(args, "s#|S:setdefault",
                              &key.dptr, &key.dsize, &defvalue))
		return NULL;
        check_dbmobject_open(dp);
	val = dbm_fetch(dp->di_dbm, key);
	if (val.dptr != NULL)
		return PyString_FromStringAndSize(val.dptr, val.dsize);
	if (defvalue == NULL) {
		defvalue = PyString_FromStringAndSize(NULL, 0);
		if (defvalue == NULL)
			return NULL;
	}
	else
		Py_INCREF(defvalue);
	val.dptr = PyString_AS_STRING(defvalue);
	val.dsize = PyString_GET_SIZE(defvalue);
	if (dbm_store(dp->di_dbm, key, val, DBM_INSERT) < 0) {
		dbm_clearerr(dp->di_dbm);
Fred Drake's avatar
Fred Drake committed
258
		PyErr_SetString(DbmError, "cannot add item to database");
259 260 261 262 263
		return NULL;
	}
	return defvalue;
}

264
static PyMethodDef dbm_methods[] = {
265 266 267 268 269 270 271 272 273 274 275 276 277
	{"close",	(PyCFunction)dbm__close,	METH_VARARGS,
	 "close()\nClose the database."},
	{"keys",	(PyCFunction)dbm_keys,		METH_VARARGS,
	 "keys() -> list\nReturn a list of all keys in the database."},
	{"has_key",	(PyCFunction)dbm_has_key,	METH_VARARGS,
	 "has_key(key} -> boolean\nReturn true iff key is in the database."},
	{"get",		(PyCFunction)dbm_get,		METH_VARARGS,
	 "get(key[, default]) -> value\n"
	 "Return the value for key if present, otherwise default."},
	{"setdefault",	(PyCFunction)dbm_setdefault,	METH_VARARGS,
	 "setdefault(key[, default]) -> value\n"
	 "Return the value for key if present, otherwise default.  If key\n"
	 "is not in the database, it is inserted with default as the value."},
278 279 280
	{NULL,		NULL}		/* sentinel */
};

281
static PyObject *
282
dbm_getattr(dbmobject *dp, char *name)
283
{
284
	return Py_FindMethod(dbm_methods, (PyObject *)dp, name);
285 286
}

287
static PyTypeObject Dbmtype = {
288
	PyObject_HEAD_INIT(NULL)
289
	0,
290
	"dbm",
291 292
	sizeof(dbmobject),
	0,
293 294
	(destructor)dbm_dealloc,  /*tp_dealloc*/
	0,			  /*tp_print*/
295
	(getattrfunc)dbm_getattr, /*tp_getattr*/
296 297 298 299 300 301
	0,			  /*tp_setattr*/
	0,			  /*tp_compare*/
	0,			  /*tp_repr*/
	0,			  /*tp_as_number*/
	0,			  /*tp_as_sequence*/
	&dbm_as_mapping,	  /*tp_as_mapping*/
302 303 304 305
};

/* ----------------------------------------------------------------- */

306
static PyObject *
307
dbmopen(PyObject *self, PyObject *args)
308
{
309 310 311 312
	char *name;
	char *flags = "r";
	int iflags;
	int mode = 0666;
313

314
        if ( !PyArg_ParseTuple(args, "s|si:open", &name, &flags, &mode) )
315
		return NULL;
316
	if ( strcmp(flags, "r") == 0 )
317
		iflags = O_RDONLY;
318
	else if ( strcmp(flags, "w") == 0 )
319
		iflags = O_RDWR;
320
	else if ( strcmp(flags, "rw") == 0 ) /* B/W compat */
321
		iflags = O_RDWR|O_CREAT; 
322
	else if ( strcmp(flags, "c") == 0 )
323
		iflags = O_RDWR|O_CREAT;
324
	else if ( strcmp(flags, "n") == 0 )
325
		iflags = O_RDWR|O_CREAT|O_TRUNC;
326
	else {
327
		PyErr_SetString(DbmError,
Fred Drake's avatar
Fred Drake committed
328
				"arg 2 to open should be 'r', 'w', 'c', or 'n'");
329
		return NULL;
330 331 332 333
	}
        return newdbmobject(name, iflags, mode);
}

334
static PyMethodDef dbmmodule_methods[] = {
335 336 337
	{ "open", (PyCFunction)dbmopen, METH_VARARGS,
	  "open(path[, flag[, mode]]) -> mapping\n"
	  "Return a database object."},
338
	{ 0, 0 },
339 340
};

341
DL_EXPORT(void)
342
initdbm(void) {
343
	PyObject *m, *d, *s;
344

345
	Dbmtype.ob_type = &PyType_Type;
346 347
	m = Py_InitModule("dbm", dbmmodule_methods);
	d = PyModule_GetDict(m);
348 349 350 351 352 353 354
	if (DbmError == NULL)
		DbmError = PyErr_NewException("dbm.error", NULL, NULL);
	s = PyString_FromString(which_dbm);
	if (s != NULL) {
		PyDict_SetItemString(d, "library", s);
		Py_DECREF(s);
	}
355 356
	if (DbmError != NULL)
		PyDict_SetItemString(d, "error", DbmError);
357
}