1 files changed, 83 insertions, 35 deletions

diff --git a/Lib/test/test_shelve.py b/Lib/test/test_shelve.py
index df56625a599..3e73f52b87b 100644
--- a/Lib/test/test_shelve.py
+++ b/Lib/test/test_shelve.py
@@ -1,14 +1,51 @@
-import os
 import unittest
 import shelve
 import glob
-from test import test_support
+from test import support
+from collections import MutableMapping
+from test.test_dbm import dbm_iterator
+
+def L1(s):
+    return s.decode("latin-1")
+
+class byteskeydict(MutableMapping):
+    "Mapping that supports bytes keys"
+
+    def __init__(self):
+        self.d = {}
+
+    def __getitem__(self, key):
+        return self.d[L1(key)]
+
+    def __setitem__(self, key, value):
+        self.d[L1(key)] = value
+
+    def __delitem__(self, key):
+        del self.d[L1(key)]
+
+    def __len__(self):
+        return len(self.d)
+
+    def iterkeys(self):
+        for k in self.d.keys():
+            yield k.encode("latin-1")
+
+    __iter__ = iterkeys
+
+    def keys(self):
+        return list(self.iterkeys())
+
+    def copy(self):
+        return byteskeydict(self.d)
 
-test_support.import_module('anydbm', deprecated=True)
 
 class TestCase(unittest.TestCase):
 
-    fn = "shelftemp" + os.extsep + "db"
+    fn = "shelftemp.db"
+
+    def tearDown(self):
+        for f in glob.glob(self.fn+"*"):
+            support.unlink(f)
 
     def test_close(self):
         d1 = {}
@@ -26,52 +63,47 @@ class TestCase(unittest.TestCase):
             self.fail('Closed shelf should not find a key')
 
     def test_ascii_file_shelf(self):
+        s = shelve.open(self.fn, protocol=0)
         try:
-            s = shelve.open(self.fn, protocol=0)
             s['key1'] = (1,2,3,4)
             self.assertEqual(s['key1'], (1,2,3,4))
-            s.close()
         finally:
-            for f in glob.glob(self.fn+"*"):
-                os.unlink(f)
+            s.close()
 
     def test_binary_file_shelf(self):
+        s = shelve.open(self.fn, protocol=1)
         try:
-            s = shelve.open(self.fn, protocol=1)
             s['key1'] = (1,2,3,4)
             self.assertEqual(s['key1'], (1,2,3,4))
-            s.close()
         finally:
-            for f in glob.glob(self.fn+"*"):
-                os.unlink(f)
+            s.close()
 
     def test_proto2_file_shelf(self):
+        s = shelve.open(self.fn, protocol=2)
         try:
-            s = shelve.open(self.fn, protocol=2)
             s['key1'] = (1,2,3,4)
             self.assertEqual(s['key1'], (1,2,3,4))
-            s.close()
         finally:
-            for f in glob.glob(self.fn+"*"):
-                os.unlink(f)
+            s.close()
 
     def test_in_memory_shelf(self):
-        d1 = {}
+        d1 = byteskeydict()
         s = shelve.Shelf(d1, protocol=0)
         s['key1'] = (1,2,3,4)
         self.assertEqual(s['key1'], (1,2,3,4))
         s.close()
-        d2 = {}
+        d2 = byteskeydict()
         s = shelve.Shelf(d2, protocol=1)
         s['key1'] = (1,2,3,4)
         self.assertEqual(s['key1'], (1,2,3,4))
         s.close()
 
         self.assertEqual(len(d1), 1)
-        self.assertNotEqual(d1, d2)
+        self.assertEqual(len(d2), 1)
+        self.assertNotEqual(d1.items(), d2.items())
 
     def test_mutable_entry(self):
-        d1 = {}
+        d1 = byteskeydict()
         s = shelve.Shelf(d1, protocol=2, writeback=False)
         s['key1'] = [1,2,3,4]
         self.assertEqual(s['key1'], [1,2,3,4])
@@ -79,7 +111,7 @@ class TestCase(unittest.TestCase):
         self.assertEqual(s['key1'], [1,2,3,4])
         s.close()
 
-        d2 = {}
+        d2 = byteskeydict()
         s = shelve.Shelf(d2, protocol=2, writeback=True)
         s['key1'] = [1,2,3,4]
         self.assertEqual(s['key1'], [1,2,3,4])
@@ -90,15 +122,30 @@ class TestCase(unittest.TestCase):
         self.assertEqual(len(d1), 1)
         self.assertEqual(len(d2), 1)
 
+    def test_keyencoding(self):
+        d = {}
+        key = 'Pöp'
+        # the default keyencoding is utf-8
+        shelve.Shelf(d)[key] = [1]
+        self.assertIn(key.encode('utf-8'), d)
+        # but a different one can be given
+        shelve.Shelf(d, keyencoding='latin1')[key] = [1]
+        self.assertIn(key.encode('latin1'), d)
+        # with all consequences
+        s = shelve.Shelf(d, keyencoding='ascii')
+        self.assertRaises(UnicodeEncodeError, s.__setitem__, key, [1])
+
     def test_writeback_also_writes_immediately(self):
         # Issue 5754
         d = {}
+        key = 'key'
+        encodedkey = key.encode('utf-8')
         s = shelve.Shelf(d, writeback=True)
-        s['key'] = [1]
-        p1 = d['key']  # Will give a KeyError if backing store not updated
+        s[key] = [1]
+        p1 = d[encodedkey]  # Will give a KeyError if backing store not updated
         s['key'].append(2)
         s.close()
-        p2 = d['key']
+        p2 = d[encodedkey]
         self.assertNotEqual(p1, p2)  # Write creates new object in store
 
 
@@ -115,7 +162,7 @@ class TestShelveBase(mapping_tests.BasicTestMappingProtocol):
         return {"key1":"value1", "key2":2, "key3":(1,2,3)}
     def _empty_mapping(self):
         if self._in_mem:
-            x= shelve.Shelf({}, **self._args)
+            x= shelve.Shelf(byteskeydict(), **self._args)
         else:
             self.counter+=1
             x= shelve.open(self.fn+str(self.counter), **self._args)
@@ -127,7 +174,7 @@ class TestShelveBase(mapping_tests.BasicTestMappingProtocol):
         self._db = []
         if not self._in_mem:
             for f in glob.glob(self.fn+"*"):
-                test_support.unlink(f)
+                support.unlink(f)
 
 class TestAsciiFileShelve(TestShelveBase):
     _args={'protocol':0}
@@ -149,15 +196,16 @@ class TestProto2MemShelve(TestShelveBase):
     _in_mem = True
 
 def test_main():
-    test_support.run_unittest(
-        TestAsciiFileShelve,
-        TestBinaryFileShelve,
-        TestProto2FileShelve,
-        TestAsciiMemShelve,
-        TestBinaryMemShelve,
-        TestProto2MemShelve,
-        TestCase
-    )
+    for module in dbm_iterator():
+        support.run_unittest(
+            TestAsciiFileShelve,
+            TestBinaryFileShelve,
+            TestProto2FileShelve,
+            TestAsciiMemShelve,
+            TestBinaryMemShelve,
+            TestProto2MemShelve,
+            TestCase
+        )
 
 if __name__ == "__main__":
     test_main()