Commit a54a8bab authored by Luke Plant's avatar Luke Plant
Browse files

Fixed #17776 - DoesNotExist is not picklable

Thanks to ambv for the report
parent f08fa5b5
Loading
Loading
Loading
Loading
+33 −6
Original line number Diff line number Diff line
@@ -62,11 +62,13 @@ class ModelBase(type):
            new_class.add_to_class('DoesNotExist', subclass_exception(b'DoesNotExist',
                    tuple(x.DoesNotExist
                          for x in parents if hasattr(x, '_meta') and not x._meta.abstract)
                                    or (ObjectDoesNotExist,), module))
                    or (ObjectDoesNotExist,),
                    module, attached_to=new_class))
            new_class.add_to_class('MultipleObjectsReturned', subclass_exception(b'MultipleObjectsReturned',
                    tuple(x.MultipleObjectsReturned
                          for x in parents if hasattr(x, '_meta') and not x._meta.abstract)
                                    or (MultipleObjectsReturned,), module))
                    or (MultipleObjectsReturned,),
                    module, attached_to=new_class))
            if base_meta and not base_meta.abstract:
                # Non-abstract child classes inherit some attributes from their
                # non-abstract parent (unless an ABC comes before it in the
@@ -934,5 +936,30 @@ def model_unpickle(model, attrs, factory):
    return cls.__new__(cls)
model_unpickle.__safe_for_unpickle__ = True

def subclass_exception(name, parents, module):
    return type(name, parents, {'__module__': module})
def subclass_exception(name, parents, module, attached_to=None):
    """
    Create exception subclass.

    If 'attached_to' is supplied, the exception will be created in a way that
    allows it to be pickled, assuming the returned exception class will be added
    as an attribute to the 'attached_to' class.
    """
    class_dict = {'__module__': module}
    if attached_to is not None:
        def __reduce__(self):
            # Exceptions are special - they've got state that isn't
            # in self.__dict__. We assume it is all in self.args.
            return (unpickle_inner_exception, (attached_to, name), self.args)

        def __setstate__(self, args):
            self.args = args

        class_dict['__reduce__'] = __reduce__
        class_dict['__setstate__'] = __setstate__

    return type(name, parents, class_dict)

def unpickle_inner_exception(klass, exception_name):
    # Get the exception class from the class it is attached to:
    exception = getattr(klass, exception_name)
    return exception.__new__(exception)
+10 −0
Original line number Diff line number Diff line
@@ -36,3 +36,13 @@ class PickleabilityTestCase(TestCase):

    def test_membermethod_as_default(self):
        self.assert_pickles(Happening.objects.filter(number4=1))

    def test_doesnotexist_exception(self):
        # Ticket #17776
        original = Event.DoesNotExist("Doesn't exist")
        unpickled = pickle.loads(pickle.dumps(original))

        # Exceptions are not equal to equivalent instances of themselves, so
        # can't just use assertEqual(original, unpickled)
        self.assertEqual(original.__class__, unpickled.__class__)
        self.assertEqual(original.args, unpickled.args)