1edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepimport copy_reg
2edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepimport unittest
3edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
4edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepfrom test import test_support
5edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepfrom test.pickletester import ExtensionSaver
6edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
7edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepclass C:
8edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    pass
9edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
10edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
11edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepclass WithoutSlots(object):
12edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    pass
13edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
14edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepclass WithWeakref(object):
15edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    __slots__ = ('__weakref__',)
16edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
17edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepclass WithPrivate(object):
18edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    __slots__ = ('__spam',)
19edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
20edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepclass WithSingleString(object):
21edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    __slots__ = 'spam'
22edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
23edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepclass WithInherited(WithSingleString):
24edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    __slots__ = ('eggs',)
25edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
26edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
27edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepclass CopyRegTestCase(unittest.TestCase):
28edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
29edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def test_class(self):
30edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        self.assertRaises(TypeError, copy_reg.pickle,
31edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                          C, None, None)
32edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
33edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def test_noncallable_reduce(self):
34edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        self.assertRaises(TypeError, copy_reg.pickle,
35edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                          type(1), "not a callable")
36edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
37edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def test_noncallable_constructor(self):
38edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        self.assertRaises(TypeError, copy_reg.pickle,
39edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                          type(1), int, "not a callable")
40edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
41edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def test_bool(self):
42edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        import copy
43edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        self.assertEqual(True, copy.copy(True))
44edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
45edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def test_extension_registry(self):
46edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        mod, func, code = 'junk1 ', ' junk2', 0xabcd
47edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        e = ExtensionSaver(code)
48edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        try:
49edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            # Shouldn't be in registry now.
50edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertRaises(ValueError, copy_reg.remove_extension,
51edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                              mod, func, code)
52edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            copy_reg.add_extension(mod, func, code)
53edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            # Should be in the registry.
54edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertTrue(copy_reg._extension_registry[mod, func] == code)
55edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertTrue(copy_reg._inverted_registry[code] == (mod, func))
56edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            # Shouldn't be in the cache.
57edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertNotIn(code, copy_reg._extension_cache)
58edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            # Redundant registration should be OK.
59edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            copy_reg.add_extension(mod, func, code)  # shouldn't blow up
60edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            # Conflicting code.
61edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertRaises(ValueError, copy_reg.add_extension,
62edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                              mod, func, code + 1)
63edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertRaises(ValueError, copy_reg.remove_extension,
64edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                              mod, func, code + 1)
65edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            # Conflicting module name.
66edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertRaises(ValueError, copy_reg.add_extension,
67edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                              mod[1:], func, code )
68edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertRaises(ValueError, copy_reg.remove_extension,
69edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                              mod[1:], func, code )
70edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            # Conflicting function name.
71edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertRaises(ValueError, copy_reg.add_extension,
72edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                              mod, func[1:], code)
73edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertRaises(ValueError, copy_reg.remove_extension,
74edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                              mod, func[1:], code)
75edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            # Can't remove one that isn't registered at all.
76edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            if code + 1 not in copy_reg._inverted_registry:
77edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                self.assertRaises(ValueError, copy_reg.remove_extension,
78edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                                  mod[1:], func[1:], code + 1)
79edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
80edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        finally:
81edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            e.restore()
82edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
83edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        # Shouldn't be there anymore.
84edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        self.assertNotIn((mod, func), copy_reg._extension_registry)
85edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        # The code *may* be in copy_reg._extension_registry, though, if
86edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        # we happened to pick on a registered code.  So don't check for
87edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        # that.
88edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
89edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        # Check valid codes at the limits.
90edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        for code in 1, 0x7fffffff:
91edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            e = ExtensionSaver(code)
92edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            try:
93edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                copy_reg.add_extension(mod, func, code)
94edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                copy_reg.remove_extension(mod, func, code)
95edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            finally:
96edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                e.restore()
97edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
98edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        # Ensure invalid codes blow up.
99edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        for code in -1, 0, 0x80000000L:
100edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep            self.assertRaises(ValueError, copy_reg.add_extension,
101edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep                              mod, func, code)
102edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
103edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    def test_slotnames(self):
104edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        self.assertEqual(copy_reg._slotnames(WithoutSlots), [])
105edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        self.assertEqual(copy_reg._slotnames(WithWeakref), [])
106edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        expected = ['_WithPrivate__spam']
107edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        self.assertEqual(copy_reg._slotnames(WithPrivate), expected)
108edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        self.assertEqual(copy_reg._slotnames(WithSingleString), ['spam'])
109edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        expected = ['eggs', 'spam']
110edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        expected.sort()
111edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        result = copy_reg._slotnames(WithInherited)
112edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        result.sort()
113edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep        self.assertEqual(result, expected)
114edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
115edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
116edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepdef test_main():
117edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    test_support.run_unittest(CopyRegTestCase)
118edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
119edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep
120edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoepif __name__ == "__main__":
121edbb763a2b63074cd468a5d33a17908b2cc0654Jeff Vander Stoep    test_main()
122