Close #19282: Native context management in dbm

4 files changed, 43 insertions, 0 deletions

diff --git a/Lib/dbm/dumb.py b/Lib/dbm/dumb.py
index 9ac7852..ba6a20d 100644
--- a/Lib/dbm/dumb.py
+++ b/Lib/dbm/dumb.py
@@ -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.
diff --git a/Lib/test/test_dbm_dumb.py b/Lib/test/test_dbm_dumb.py
index 208bc4c..9fd8cde 100644
--- a/Lib/test/test_dbm_dumb.py
+++ b/Lib/test/test_dbm_dumb.py
@@ -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()
 
diff --git a/Lib/test/test_dbm_gnu.py b/Lib/test/test_dbm_gnu.py
index 4fb66c5..a7808f5 100755
--- a/Lib/test/test_dbm_gnu.py
+++ b/Lib/test/test_dbm_gnu.py
@@ -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()
diff --git a/Lib/test/test_dbm_ndbm.py b/Lib/test/test_dbm_ndbm.py
index b57e1f0..2291561 100755
--- a/Lib/test/test_dbm_ndbm.py
+++ b/Lib/test/test_dbm_ndbm.py
@@ -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()