183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh"""Unit tests for collections.defaultdict."""
283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport os
483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport copy
583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport tempfile
683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport unittest
783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom test import test_support
883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom collections import defaultdict
1083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
1183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef foobar():
1283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    return list
1383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
1483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehclass TestDefaultDict(unittest.TestCase):
1583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
1683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def test_basic(self):
1783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1 = defaultdict()
1883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d1.default_factory, None)
1983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1.default_factory = list
2083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1[12].append(42)
2183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d1, {12: [42]})
2283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1[12].append(24)
2383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d1, {12: [42, 24]})
2483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1[13]
2583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1[14]
2683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d1, {12: [42, 24], 13: [], 14: []})
2783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertTrue(d1[12] is not d1[13] is not d1[14])
2883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d2 = defaultdict(list, foo=1, bar=2)
2983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2.default_factory, list)
3083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2, {"foo": 1, "bar": 2})
3183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2["foo"], 1)
3283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2["bar"], 2)
3383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2[42], [])
3483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertIn("foo", d2)
3583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertIn("foo", d2.keys())
3683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertIn("bar", d2)
3783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertIn("bar", d2.keys())
3883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertIn(42, d2)
3983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertIn(42, d2.keys())
4083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertNotIn(12, d2)
4183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertNotIn(12, d2.keys())
4283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d2.default_factory = None
4383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2.default_factory, None)
4483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
4583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            d2[15]
4683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        except KeyError, err:
4783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.assertEqual(err.args, (15,))
4883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
4983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.fail("d2[15] didn't raise KeyError")
5083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertRaises(TypeError, defaultdict, 1)
5183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
5283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def test_missing(self):
5383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1 = defaultdict()
5483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertRaises(KeyError, d1.__missing__, 42)
5583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1.default_factory = list
5683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d1.__missing__(42), [])
5783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
5883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def test_repr(self):
5983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1 = defaultdict()
6083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d1.default_factory, None)
6183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(repr(d1), "defaultdict(None, {})")
6283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(eval(repr(d1)), d1)
6383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1[11] = 41
6483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(repr(d1), "defaultdict(None, {11: 41})")
6583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d2 = defaultdict(int)
6683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2.default_factory, int)
6783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d2[12] = 42
6883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(repr(d2), "defaultdict(<type 'int'>, {12: 42})")
6983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def foo(): return 43
7083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d3 = defaultdict(foo)
7183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertTrue(d3.default_factory is foo)
7283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d3[13]
7383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(repr(d3), "defaultdict(%s, {13: 43})" % repr(foo))
7483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
7583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def test_print(self):
7683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1 = defaultdict()
7783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def foo(): return 42
7883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d2 = defaultdict(foo, {1: 2})
7983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # NOTE: We can't use tempfile.[Named]TemporaryFile since this
8083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # code must exercise the tp_print C code, which only gets
8183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # invoked for *real* files.
8283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        tfn = tempfile.mktemp()
8383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
8483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            f = open(tfn, "w+")
8583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            try:
8683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                print >>f, d1
8783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                print >>f, d2
8883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                f.seek(0)
8983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.assertEqual(f.readline(), repr(d1) + "\n")
9083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.assertEqual(f.readline(), repr(d2) + "\n")
9183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            finally:
9283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                f.close()
9383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        finally:
9483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            os.remove(tfn)
9583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
9683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def test_copy(self):
9783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1 = defaultdict()
9883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d2 = d1.copy()
9983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(type(d2), defaultdict)
10083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2.default_factory, None)
10183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2, {})
10283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1.default_factory = list
10383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d3 = d1.copy()
10483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(type(d3), defaultdict)
10583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d3.default_factory, list)
10683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d3, {})
10783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1[42]
10883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d4 = d1.copy()
10983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(type(d4), defaultdict)
11083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d4.default_factory, list)
11183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d4, {42: []})
11283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d4[12]
11383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d4, {42: [], 12: []})
11483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
11583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # Issue 6637: Copy fails for empty default dict
11683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d = defaultdict()
11783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d['a'] = 42
11883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        e = d.copy()
11983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(e['a'], 42)
12083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
12183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def test_shallow_copy(self):
12283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1 = defaultdict(foobar, {1: 1})
12383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d2 = copy.copy(d1)
12483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2.default_factory, foobar)
12583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2, d1)
12683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1.default_factory = list
12783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d2 = copy.copy(d1)
12883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2.default_factory, list)
12983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2, d1)
13083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
13183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def test_deep_copy(self):
13283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1 = defaultdict(foobar, {1: [1]})
13383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d2 = copy.deepcopy(d1)
13483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2.default_factory, foobar)
13583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2, d1)
13683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertTrue(d1[1] is not d2[1])
13783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1.default_factory = list
13883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d2 = copy.deepcopy(d1)
13983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2.default_factory, list)
14083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertEqual(d2, d1)
14183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
14283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def test_keyerror_without_factory(self):
14383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d1 = defaultdict()
14483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
14583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            d1[(1,)]
14683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        except KeyError, err:
14783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.assertEqual(err.args[0], (1,))
14883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
14983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.fail("expected KeyError")
15083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
15183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def test_recursive_repr(self):
15283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # Issue2045: stack overflow when default_factory is a bound method
15383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        class sub(defaultdict):
15483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            def __init__(self):
15583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.default_factory = self._factory
15683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            def _factory(self):
15783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                return []
15883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d = sub()
15983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertTrue(repr(d).startswith(
16083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            "defaultdict(<bound method sub._factory of defaultdict(..."))
16183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
16283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # NOTE: printing a subclass of a builtin type does not call its
16383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # tp_print slot. So this part is essentially the same test as above.
16483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        tfn = tempfile.mktemp()
16583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
16683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            f = open(tfn, "w+")
16783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            try:
16883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                print >>f, d
16983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            finally:
17083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                f.close()
17183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        finally:
17283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            os.remove(tfn)
17383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
17483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def test_callable_arg(self):
17583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.assertRaises(TypeError, defaultdict, {})
17683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
17783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef test_main():
17883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    test_support.run_unittest(TestDefaultDict)
17983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
18083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehif __name__ == "__main__":
18183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    test_main()
182