import unittest
import pickle
import cPickle
import os

class A(Exception):
    """Simple extension."""

class B(Exception):
    """Extension with values."""
    def __init__(self, foo):
        self.foo = foo
        Exception.__init__(self, foo)

class C(Exception):
    """Extension with values, args not set."""
    def __init__(self, foo):
        self.foo = foo

class Test(unittest.TestCase):

    def test_exception_pickle(self):

        # Build an OSError exception with errno set.
        try:
            os.unlink('some_file_doesnt_exist')
        except OSError, e:
            pass

        values = [
            # Some standard expetions.
            ValueError(7),
            e,
            # Custom exceptions.
            A(),
            A(1, 2, 3),
            B(3),
            C(4),
        ]

        def compare(exc_obj, unpickled_exc_obj):
            if hasattr(exc_obj, 'args'):
                self.assertEqual(exc_obj.args, unpickled_exc_obj.args)
            else:
                self.assertFalse(hasattr(unpickled_exc_obj, 'args'))
            self.assertEqual(exc_obj.__dict__, unpickled_exc_obj.__dict__)

        for proto in xrange(pickle.HIGHEST_PROTOCOL+1):
            for module in (pickle, cPickle):
                for value in values:
                    pickle_str = module.dumps(value, proto)
                    unpickled = module.loads(pickle_str)
                    compare(value, unpickled)

        # Test old pickle exceptions from 2.4 (pre-new-style exceptions).
        old_pickle_strings = {
            0: ["(iexceptions\nValueError\np0\n(dp1\nS'args'\np2\n(I7\ntp3\nsb.",
                "(iexceptions\nOSError\np0\n(dp1\nS'errno'\np2\nI2\nsS'args'\np3\n(I2\nS'No such file or directory'\np4\ntp5\nsS'strerror'\np6\ng4\nsS'filename'\np7\nS'some_file_doesnt_exist'\np8\nsb.",
                "(i__main__\nA\np0\n(dp1\nS'args'\np2\n(tsb.",
                "(i__main__\nA\np0\n(dp1\nS'args'\np2\n(I1\nI2\nI3\ntp3\nsb.",
                "(i__main__\nB\np0\n(dp1\nS'foo'\np2\nI3\nsS'args'\np3\n(I3\ntp4\nsb.",
                "(i__main__\nC\np0\n(dp1\nS'foo'\np2\nI4\nsb.",
               ],
            1: ['(cexceptions\nValueError\nq\x00oq\x01}q\x02U\x04argsq\x03(K\x07tq\x04sb.',
                '(cexceptions\nOSError\nq\x00oq\x01}q\x02(U\x05errnoq\x03K\x02U\x04argsq\x04(K\x02U\x19No such file or directoryq\x05tq\x06U\x08strerrorq\x07h\x05U\x08filenameq\x08U\x16some_file_doesnt_existq\tub.',
                '(c__main__\nA\nq\x00oq\x01}q\x02U\x04argsq\x03)sb.',
                '(c__main__\nA\nq\x00oq\x01}q\x02U\x04argsq\x03(K\x01K\x02K\x03tq\x04sb.',
                '(c__main__\nB\nq\x00oq\x01}q\x02(U\x03fooq\x03K\x03U\x04argsq\x04(K\x03tq\x05ub.',
                '(c__main__\nC\nq\x00oq\x01}q\x02U\x03fooq\x03K\x04sb.',
               ],
            2: ['\x80\x02(cexceptions\nValueError\nq\x00oq\x01}q\x02U\x04argsq\x03K\x07\x85q\x04sb.',
                '\x80\x02(cexceptions\nOSError\nq\x00oq\x01}q\x02(U\x05errnoq\x03K\x02U\x04argsq\x04K\x02U\x19No such file or directoryq\x05\x86q\x06U\x08strerrorq\x07h\x05U\x08filenameq\x08U\x16some_file_doesnt_existq\tub.',
                '\x80\x02(c__main__\nA\nq\x00oq\x01}q\x02U\x04argsq\x03)sb.',
                '\x80\x02(c__main__\nA\nq\x00oq\x01}q\x02U\x04argsq\x03K\x01K\x02K\x03\x87q\x04sb.',
                '\x80\x02(c__main__\nB\nq\x00oq\x01}q\x02(U\x03fooq\x03K\x03U\x04argsq\x04K\x03\x85q\x05ub.',
                '\x80\x02(c__main__\nC\nq\x00oq\x01}q\x02U\x03fooq\x03K\x04sb.',
               ],
        }

        for proto, pickle_strings in old_pickle_strings.items():
            for module in (pickle, cPickle):
                for value, pickle_str in zip(values, pickle_strings):
                    #print repr((proto, pickle_str, value))
                    unpickled = module.loads(pickle_str)
                    compare(value, unpickled)

if __name__ == '__main__':
    unittest.main()