import unittest, pickle as cPickle
from test import support
pickle = support.import_fresh_module('pickle', blocked=('_pickle',))

class EnumItem(int):
    def __new__(cls, val, an_enum):
        res = int.__new__(cls, val)
        res.an_enum = an_enum
        return res
    def __getnewargs__(self):
        return (int(self), self.an_enum)
    def __getstate__(self):
        return int(self), self.an_enum
    def __setstate__(self, state):
        ival,an_enum = state
        assert ival==int(self)
        self.an_enum = an_enum
        
class Enum(object):
    def __init__(self):
        self.a = EnumItem(1, self)

class TestEnumPickling(unittest.TestCase):
    def setUp(self):
        E = Enum()
        self.orig = [
            set([E.a]),
            set([E.a]),
            ]

    def check(self, picklemod, protocol):
        dupe = picklemod.loads(picklemod.dumps(self.orig, protocol=protocol))
        self.assertEqual(self.orig, dupe)
        
    def tests_that_succeed(self):
        for picklemod in (pickle, cPickle):
            for proto in (0,1,):
                self.check(picklemod, proto)

    def test_pickle_protocol2(self):
        self.check(pickle, 2)

    def test_cPickle_protocol2(self):
        self.check(cPickle, 2)
        
    def test_pickle_protocol3(self):
        self.check(pickle, 3)

    def test_cPickle_protocol3(self):
        self.check(cPickle, 3)
        
if __name__ == '__main__':
    unittest.main()