10c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi"""Unit tests for collections.defaultdict."""
20c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
30c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport os
40c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport copy
50c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport tempfile
60c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport unittest
70c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yifrom test import test_support
80c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
90c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yifrom collections import defaultdict
100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef foobar():
120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    return list
130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass TestDefaultDict(unittest.TestCase):
150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def test_basic(self):
170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1 = defaultdict()
180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d1.default_factory, None)
190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1.default_factory = list
200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1[12].append(42)
210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d1, {12: [42]})
220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1[12].append(24)
230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d1, {12: [42, 24]})
240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1[13]
250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1[14]
260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d1, {12: [42, 24], 13: [], 14: []})
270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertTrue(d1[12] is not d1[13] is not d1[14])
280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d2 = defaultdict(list, foo=1, bar=2)
290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2.default_factory, list)
300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2, {"foo": 1, "bar": 2})
310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2["foo"], 1)
320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2["bar"], 2)
330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2[42], [])
340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertIn("foo", d2)
350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertIn("foo", d2.keys())
360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertIn("bar", d2)
370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertIn("bar", d2.keys())
380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertIn(42, d2)
390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertIn(42, d2.keys())
400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertNotIn(12, d2)
410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertNotIn(12, d2.keys())
420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d2.default_factory = None
430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2.default_factory, None)
440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        try:
450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            d2[15]
460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        except KeyError, err:
470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            self.assertEqual(err.args, (15,))
480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        else:
490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            self.fail("d2[15] didn't raise KeyError")
500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertRaises(TypeError, defaultdict, 1)
510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def test_missing(self):
530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1 = defaultdict()
540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertRaises(KeyError, d1.__missing__, 42)
550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1.default_factory = list
560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d1.__missing__(42), [])
570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def test_repr(self):
590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1 = defaultdict()
600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d1.default_factory, None)
610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(repr(d1), "defaultdict(None, {})")
620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(eval(repr(d1)), d1)
630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1[11] = 41
640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(repr(d1), "defaultdict(None, {11: 41})")
650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d2 = defaultdict(int)
660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2.default_factory, int)
670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d2[12] = 42
680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(repr(d2), "defaultdict(<type 'int'>, {12: 42})")
690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        def foo(): return 43
700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d3 = defaultdict(foo)
710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertTrue(d3.default_factory is foo)
720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d3[13]
730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(repr(d3), "defaultdict(%s, {13: 43})" % repr(foo))
740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def test_print(self):
760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1 = defaultdict()
770c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        def foo(): return 42
780c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d2 = defaultdict(foo, {1: 2})
790c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        # NOTE: We can't use tempfile.[Named]TemporaryFile since this
800c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        # code must exercise the tp_print C code, which only gets
810c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        # invoked for *real* files.
820c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        tfn = tempfile.mktemp()
830c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        try:
840c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            f = open(tfn, "w+")
850c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            try:
860c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                print >>f, d1
870c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                print >>f, d2
880c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                f.seek(0)
890c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                self.assertEqual(f.readline(), repr(d1) + "\n")
900c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                self.assertEqual(f.readline(), repr(d2) + "\n")
910c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            finally:
920c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                f.close()
930c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        finally:
940c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            os.remove(tfn)
950c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
960c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def test_copy(self):
970c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1 = defaultdict()
980c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d2 = d1.copy()
990c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(type(d2), defaultdict)
1000c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2.default_factory, None)
1010c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2, {})
1020c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1.default_factory = list
1030c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d3 = d1.copy()
1040c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(type(d3), defaultdict)
1050c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d3.default_factory, list)
1060c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d3, {})
1070c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1[42]
1080c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d4 = d1.copy()
1090c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(type(d4), defaultdict)
1100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d4.default_factory, list)
1110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d4, {42: []})
1120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d4[12]
1130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d4, {42: [], 12: []})
1140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        # Issue 6637: Copy fails for empty default dict
1160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d = defaultdict()
1170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d['a'] = 42
1180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        e = d.copy()
1190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(e['a'], 42)
1200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def test_shallow_copy(self):
1220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1 = defaultdict(foobar, {1: 1})
1230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d2 = copy.copy(d1)
1240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2.default_factory, foobar)
1250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2, d1)
1260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1.default_factory = list
1270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d2 = copy.copy(d1)
1280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2.default_factory, list)
1290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2, d1)
1300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def test_deep_copy(self):
1320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1 = defaultdict(foobar, {1: [1]})
1330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d2 = copy.deepcopy(d1)
1340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2.default_factory, foobar)
1350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2, d1)
1360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertTrue(d1[1] is not d2[1])
1370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1.default_factory = list
1380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d2 = copy.deepcopy(d1)
1390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2.default_factory, list)
1400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertEqual(d2, d1)
1410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def test_keyerror_without_factory(self):
1430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d1 = defaultdict()
1440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        try:
1450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            d1[(1,)]
1460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        except KeyError, err:
1470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            self.assertEqual(err.args[0], (1,))
1480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        else:
1490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            self.fail("expected KeyError")
1500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def test_recursive_repr(self):
1520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        # Issue2045: stack overflow when default_factory is a bound method
1530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        class sub(defaultdict):
1540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            def __init__(self):
1550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                self.default_factory = self._factory
1560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            def _factory(self):
1570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                return []
1580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        d = sub()
1590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertTrue(repr(d).startswith(
1600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            "defaultdict(<bound method sub._factory of defaultdict(..."))
1610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        # NOTE: printing a subclass of a builtin type does not call its
1630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        # tp_print slot. So this part is essentially the same test as above.
1640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        tfn = tempfile.mktemp()
1650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        try:
1660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            f = open(tfn, "w+")
1670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            try:
1680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                print >>f, d
1690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            finally:
1700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi                f.close()
1710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        finally:
1720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi            os.remove(tfn)
1730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    def test_callable_arg(self):
1750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi        self.assertRaises(TypeError, defaultdict, {})
1760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1770c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef test_main():
1780c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    test_support.run_unittest(TestDefaultDict)
1790c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi
1800c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiif __name__ == "__main__":
1810c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi    test_main()
182