import pickle   # python pickle module
import _pickle  # c pickle module


class A:
    """a class with no dict"""
    __slots__ = 'slotted_attribute'

    def __init__(self, slotted_attribute=None):
        self.slotted_attribute = slotted_attribute

    # a class with slots must define a __getstate__ method
    def __getstate__(self):
        return None, {'slotted_attribute': self.slotted_attribute}


class B:
    """a class with no dict but a different __getstate__ method"""
    __slots__ = 'slotted_attribute'

    def __init__(self, slotted_attribute=None):
        self.slotted_attribute = slotted_attribute

    # a class with slots must define a __getstate__ method
    def __getstate__(self):
        return {}, {'slotted_attribute': self.slotted_attribute}


class C(A):
    """subclass of a slotted class"""
    def __init__(self, non_slotted_attribute=None, slotted_attribute=None):
        self.non_slotted_attribute = non_slotted_attribute
        super(C, self).__init__(slotted_attribute)

    def __getstate__(self):
        state = {'non_slotted_attribute': self.non_slotted_attribute}
        slotstate = {'slotted_attribute': self.slotted_attribute}
        return state, slotstate


class D:
    """class with __dict__ in its slots"""
    slots = ('slotted_attribute', '__dict__')

    def __init__(self, non_slotted_attribute=None, slotted_attribute=None):
        self.non_slotted_attribute = non_slotted_attribute
        self.slotted_attribute = slotted_attribute

    def __getstate__(self):
        state = {'non_slotted_attribute': self.non_slotted_attribute}
        slotstate = {'slotted_attribute': self.slotted_attribute}
        return state, slotstate


def python_pickle_depickle(obj):
    return pickle._loads(pickle._dumps(obj))


def c_pickle_depickle(obj):
    return _pickle.loads(_pickle.dumps(obj))


if __name__ == "__main__":
    initargs = (
        {'slotted_attribute': 10},
        {'slotted_attribute': 10},
        {'slotted_attribute': 10, 'non_slotted_attribute': 20},
        {'slotted_attribute': 10, 'non_slotted_attribute': 20}
    )
    classes = (A, B, C, D)

    for cls, initargs in zip([A, B, C, D], initargs):
        for pickle_depickle in [c_pickle_depickle, python_pickle_depickle]:

            msg = 'depickling {} obj with {:30}: '.format(
                cls.__qualname__, pickle_depickle.__qualname__)
            try:
                obj = cls(**initargs)  # create a new object
                depickled_obj = pickle_depickle(obj)  # dump/load cycle

                # make sure the attributes are properly retrieved
                for k, v in initargs.items():
                    assert getattr(depickled_obj, k, v) == getattr(obj, k, v)

            except Exception as e:
                msg += 'fail: {}'.format(e)
            else:
                msg += 'ok'
            print(msg)