1import unittest
2import dbm
3import shelve
4import glob
5import pickle
6import os
7
8from test import support
9from test.support import os_helper
10from collections.abc import MutableMapping
11from test.test_dbm import dbm_iterator
12
13def L1(s):
14    return s.decode("latin-1")
15
16class byteskeydict(MutableMapping):
17    "Mapping that supports bytes keys"
18
19    def __init__(self):
20        self.d = {}
21
22    def __getitem__(self, key):
23        return self.d[L1(key)]
24
25    def __setitem__(self, key, value):
26        self.d[L1(key)] = value
27
28    def __delitem__(self, key):
29        del self.d[L1(key)]
30
31    def __len__(self):
32        return len(self.d)
33
34    def iterkeys(self):
35        for k in self.d.keys():
36            yield k.encode("latin-1")
37
38    __iter__ = iterkeys
39
40    def keys(self):
41        return list(self.iterkeys())
42
43    def copy(self):
44        return byteskeydict(self.d)
45
46
47class TestCase(unittest.TestCase):
48    dirname = os_helper.TESTFN
49    fn = os.path.join(os_helper.TESTFN, "shelftemp.db")
50
51    def test_close(self):
52        d1 = {}
53        s = shelve.Shelf(d1, protocol=2, writeback=False)
54        s['key1'] = [1,2,3,4]
55        self.assertEqual(s['key1'], [1,2,3,4])
56        self.assertEqual(len(s), 1)
57        s.close()
58        self.assertRaises(ValueError, len, s)
59        try:
60            s['key1']
61        except ValueError:
62            pass
63        else:
64            self.fail('Closed shelf should not find a key')
65
66    def test_open_template(self, filename=None, protocol=None):
67        os.mkdir(self.dirname)
68        self.addCleanup(os_helper.rmtree, self.dirname)
69        s = shelve.open(filename=filename if filename is not None else self.fn,
70                        protocol=protocol)
71        try:
72            s['key1'] = (1,2,3,4)
73            self.assertEqual(s['key1'], (1,2,3,4))
74        finally:
75            s.close()
76
77    def test_ascii_file_shelf(self):
78        self.test_open_template(protocol=0)
79
80    def test_binary_file_shelf(self):
81        self.test_open_template(protocol=1)
82
83    def test_proto2_file_shelf(self):
84        self.test_open_template(protocol=2)
85
86    def test_pathlib_path_file_shelf(self):
87        self.test_open_template(filename=os_helper.FakePath(self.fn))
88
89    def test_bytes_path_file_shelf(self):
90        self.test_open_template(filename=os.fsencode(self.fn))
91
92    def test_pathlib_bytes_path_file_shelf(self):
93        self.test_open_template(filename=os_helper.FakePath(os.fsencode(self.fn)))
94
95    def test_in_memory_shelf(self):
96        d1 = byteskeydict()
97        with shelve.Shelf(d1, protocol=0) as s:
98            s['key1'] = (1,2,3,4)
99            self.assertEqual(s['key1'], (1,2,3,4))
100        d2 = byteskeydict()
101        with shelve.Shelf(d2, protocol=1) as s:
102            s['key1'] = (1,2,3,4)
103            self.assertEqual(s['key1'], (1,2,3,4))
104
105        self.assertEqual(len(d1), 1)
106        self.assertEqual(len(d2), 1)
107        self.assertNotEqual(d1.items(), d2.items())
108
109    def test_mutable_entry(self):
110        d1 = byteskeydict()
111        with shelve.Shelf(d1, protocol=2, writeback=False) as s:
112            s['key1'] = [1,2,3,4]
113            self.assertEqual(s['key1'], [1,2,3,4])
114            s['key1'].append(5)
115            self.assertEqual(s['key1'], [1,2,3,4])
116
117        d2 = byteskeydict()
118        with shelve.Shelf(d2, protocol=2, writeback=True) as s:
119            s['key1'] = [1,2,3,4]
120            self.assertEqual(s['key1'], [1,2,3,4])
121            s['key1'].append(5)
122            self.assertEqual(s['key1'], [1,2,3,4,5])
123
124        self.assertEqual(len(d1), 1)
125        self.assertEqual(len(d2), 1)
126
127    def test_keyencoding(self):
128        d = {}
129        key = 'Pöp'
130        # the default keyencoding is utf-8
131        shelve.Shelf(d)[key] = [1]
132        self.assertIn(key.encode('utf-8'), d)
133        # but a different one can be given
134        shelve.Shelf(d, keyencoding='latin-1')[key] = [1]
135        self.assertIn(key.encode('latin-1'), d)
136        # with all consequences
137        s = shelve.Shelf(d, keyencoding='ascii')
138        self.assertRaises(UnicodeEncodeError, s.__setitem__, key, [1])
139
140    def test_writeback_also_writes_immediately(self):
141        # Issue 5754
142        d = {}
143        key = 'key'
144        encodedkey = key.encode('utf-8')
145        with shelve.Shelf(d, writeback=True) as s:
146            s[key] = [1]
147            p1 = d[encodedkey]  # Will give a KeyError if backing store not updated
148            s['key'].append(2)
149        p2 = d[encodedkey]
150        self.assertNotEqual(p1, p2)  # Write creates new object in store
151
152    def test_with(self):
153        d1 = {}
154        with shelve.Shelf(d1, protocol=2, writeback=False) as s:
155            s['key1'] = [1,2,3,4]
156            self.assertEqual(s['key1'], [1,2,3,4])
157            self.assertEqual(len(s), 1)
158        self.assertRaises(ValueError, len, s)
159        try:
160            s['key1']
161        except ValueError:
162            pass
163        else:
164            self.fail('Closed shelf should not find a key')
165
166    def test_default_protocol(self):
167        with shelve.Shelf({}) as s:
168            self.assertEqual(s._protocol, pickle.DEFAULT_PROTOCOL)
169
170
171class TestShelveBase:
172    type2test = shelve.Shelf
173
174    def _reference(self):
175        return {"key1":"value1", "key2":2, "key3":(1,2,3)}
176
177
178class TestShelveInMemBase(TestShelveBase):
179    def _empty_mapping(self):
180        return shelve.Shelf(byteskeydict(), **self._args)
181
182
183class TestShelveFileBase(TestShelveBase):
184    counter = 0
185
186    def _empty_mapping(self):
187        self.counter += 1
188        x = shelve.open(self.base_path + str(self.counter), **self._args)
189        self.addCleanup(x.close)
190        return x
191
192    def setUp(self):
193        dirname = os_helper.TESTFN
194        os.mkdir(dirname)
195        self.addCleanup(os_helper.rmtree, dirname)
196        self.base_path = os.path.join(dirname, "shelftemp.db")
197        self.addCleanup(setattr, dbm, '_defaultmod', dbm._defaultmod)
198        dbm._defaultmod = self.dbm_mod
199
200
201from test import mapping_tests
202
203for proto in range(pickle.HIGHEST_PROTOCOL + 1):
204    bases = (TestShelveInMemBase, mapping_tests.BasicTestMappingProtocol)
205    name = f'TestProto{proto}MemShelve'
206    globals()[name] = type(name, bases,
207                           {'_args': {'protocol': proto}})
208    bases = (TestShelveFileBase, mapping_tests.BasicTestMappingProtocol)
209    for dbm_mod in dbm_iterator():
210        assert dbm_mod.__name__.startswith('dbm.')
211        suffix = dbm_mod.__name__[4:]
212        name = f'TestProto{proto}File_{suffix}Shelve'
213        globals()[name] = type(name, bases,
214                               {'dbm_mod': dbm_mod, '_args': {'protocol': proto}})
215
216
217if __name__ == "__main__":
218    unittest.main()
219