Close #19282: Native context management in dbm

This commit is contained in:
Nick Coghlan 2013-11-17 15:59:51 +10:00
parent eb8ea265ba
commit c610aba1ed
8 changed files with 101 additions and 16 deletions

View File

@ -73,33 +73,39 @@ Key and values are always stored as bytes. This means that when
strings are used they are implicitly converted to the default encoding before
being stored.
These objects also support being used in a :keyword:`with` statement, which
will automatically close them when done.
.. versionchanged:: 3.4
Added native support for the context management protocol to the objects
returned by :func:`.open`.
The following example records some hostnames and a corresponding title, and
then prints out the contents of the database::
import dbm
# Open database, creating it if necessary.
db = dbm.open('cache', 'c')
with dbm.open('cache', 'c') as db:
# Record some values
db[b'hello'] = b'there'
db['www.python.org'] = 'Python Website'
db['www.cnn.com'] = 'Cable News Network'
# Record some values
db[b'hello'] = b'there'
db['www.python.org'] = 'Python Website'
db['www.cnn.com'] = 'Cable News Network'
# Note that the keys are considered bytes now.
assert db[b'www.python.org'] == b'Python Website'
# Notice how the value is now in bytes.
assert db['www.cnn.com'] == b'Cable News Network'
# Note that the keys are considered bytes now.
assert db[b'www.python.org'] == b'Python Website'
# Notice how the value is now in bytes.
assert db['www.cnn.com'] == b'Cable News Network'
# Often-used methods of the dict interface work too.
print(db.get('python.org', b'not present'))
# Often-used methods of the dict interface work too.
print(db.get('python.org', b'not present'))
# Storing a non-string key or value will raise an exception (most
# likely a TypeError).
db['www.yahoo.com'] = 4
# Storing a non-string key or value will raise an exception (most
# likely a TypeError).
db['www.yahoo.com'] = 4
# Close when done.
db.close()
# db is automatically closed when leaving the with statement.
.. seealso::

View File

@ -236,6 +236,12 @@ class _Database(collections.MutableMapping):
if hasattr(self._os, 'chmod'):
self._os.chmod(file, self._mode)
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
def open(file, flag=None, mode=0o666):
"""Open the database file, filename, and return corresponding object.

View File

@ -184,6 +184,19 @@ class DumbDBMTestCase(unittest.TestCase):
self.assertEqual(expected, got)
f.close()
def test_context_manager(self):
with dumbdbm.open(_fname, 'c') as db:
db["dumbdbm context manager"] = "context manager"
with dumbdbm.open(_fname, 'r') as db:
self.assertEqual(list(db.keys()), [b"dumbdbm context manager"])
# This currently just raises AttributeError rather than a specific
# exception like the GNU or NDBM based implementations. See
# http://bugs.python.org/issue19385 for details.
with self.assertRaises(Exception):
db.keys()
def tearDown(self):
_delete_files()

View File

@ -81,6 +81,17 @@ class TestGdbm(unittest.TestCase):
size2 = os.path.getsize(filename)
self.assertTrue(size1 > size2 >= size0)
def test_context_manager(self):
with gdbm.open(filename, 'c') as db:
db["gdbm context manager"] = "context manager"
with gdbm.open(filename, 'r') as db:
self.assertEqual(list(db.keys()), [b"gdbm context manager"])
with self.assertRaises(gdbm.error) as cm:
db.keys()
self.assertEqual(str(cm.exception),
"GDBM object has already been closed")
if __name__ == '__main__':
unittest.main()

View File

@ -37,5 +37,18 @@ class DbmTestCase(unittest.TestCase):
except error:
self.fail()
def test_context_manager(self):
with dbm.ndbm.open(self.filename, 'c') as db:
db["ndbm context manager"] = "context manager"
with dbm.ndbm.open(self.filename, 'r') as db:
self.assertEqual(list(db.keys()), [b"ndbm context manager"])
with self.assertRaises(dbm.ndbm.error) as cm:
db.keys()
self.assertEqual(str(cm.exception),
"DBM object has already been closed")
if __name__ == '__main__':
unittest.main()

View File

@ -50,6 +50,9 @@ Core and Builtins
Library
-------
- Issue #19282: dbm.open now supports the context manager protocol. (Inital
patch by Claudiu Popa)
- Issue #8311: Added support for writing any bytes-like objects in the aifc,
sunau, and wave modules.

View File

@ -313,6 +313,21 @@ dbm_setdefault(dbmobject *dp, PyObject *args)
return defvalue;
}
static PyObject *
dbm__enter__(PyObject *self, PyObject *args)
{
Py_INCREF(self);
return self;
}
static PyObject *
dbm__exit__(PyObject *self, PyObject *args)
{
_Py_IDENTIFIER(close);
return _PyObject_CallMethodId(self, &PyId_close, NULL);
}
static PyMethodDef dbm_methods[] = {
{"close", (PyCFunction)dbm__close, METH_NOARGS,
"close()\nClose the database."},
@ -325,6 +340,8 @@ static PyMethodDef dbm_methods[] = {
"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."},
{"__enter__", dbm__enter__, METH_NOARGS, NULL},
{"__exit__", dbm__exit__, METH_VARARGS, NULL},
{NULL, NULL} /* sentinel */
};

View File

@ -425,6 +425,20 @@ dbm_sync(dbmobject *dp, PyObject *unused)
return Py_None;
}
static PyObject *
dbm__enter__(PyObject *self, PyObject *args)
{
Py_INCREF(self);
return self;
}
static PyObject *
dbm__exit__(PyObject *self, PyObject *args)
{
_Py_IDENTIFIER(close);
return _PyObject_CallMethodId(self, &PyId_close, NULL);
}
static PyMethodDef dbm_methods[] = {
{"close", (PyCFunction)dbm_close, METH_NOARGS, dbm_close__doc__},
{"keys", (PyCFunction)dbm_keys, METH_NOARGS, dbm_keys__doc__},
@ -434,6 +448,8 @@ static PyMethodDef dbm_methods[] = {
{"sync", (PyCFunction)dbm_sync, METH_NOARGS, dbm_sync__doc__},
{"get", (PyCFunction)dbm_get, METH_VARARGS, dbm_get__doc__},
{"setdefault",(PyCFunction)dbm_setdefault,METH_VARARGS, dbm_setdefault__doc__},
{"__enter__", dbm__enter__, METH_NOARGS, NULL},
{"__exit__", dbm__exit__, METH_VARARGS, NULL},
{NULL, NULL} /* sentinel */
};