diff --git a/Lib/test/test_dbm_ndbm.py b/Lib/test/test_dbm_ndbm.py --- a/Lib/test/test_dbm_ndbm.py +++ b/Lib/test/test_dbm_ndbm.py @@ -24,6 +24,7 @@ self.d[b'bytes'] = b'data' self.d['12345678910'] = '019237410982340912840198242' self.d.keys() + self.assertIn('a', self.d) self.assertIn(b'a', self.d) self.assertEqual(self.d[b'bytes'], b'data') self.d.close() diff --git a/Modules/_dbmmodule.c b/Modules/_dbmmodule.c --- a/Modules/_dbmmodule.c +++ b/Modules/_dbmmodule.c @@ -215,22 +215,22 @@ "DBM object has already been closed"); return -1; } - if (PyUnicode_Check(arg)) { + if (PyBytes_Check(arg)) { + key.dptr = PyBytes_AS_STRING(arg); + key.dsize = PyBytes_GET_SIZE(arg); + } + else if (PyUnicode_Check(arg)) { key.dptr = PyUnicode_AsUTF8AndSize(arg, &size); key.dsize = size; if (key.dptr == NULL) return -1; } - if (!PyBytes_Check(arg)) { + else { PyErr_Format(PyExc_TypeError, - "dbm key must be string, not %.100s", + "dbm key must be bytes or string, not %.100s", arg->ob_type->tp_name); return -1; } - else { - key.dptr = PyBytes_AS_STRING(arg); - key.dsize = PyBytes_GET_SIZE(arg); - } val = dbm_fetch(dp->di_dbm, key); return val.dptr != NULL; }