Commit 5459795e authored by Anssi Kääriäinen's avatar Anssi Kääriäinen
Browse files

Fixed #20289 -- pickling of dynamic models

parent 855d1305

django/db/models/base.py

+18 −6
Original line number Diff line number Diff line
@@ -451,16 +451,18 @@ class Model(six.with_metaclass(ModelBase)): 
        need to do things manually, as they're dynamically created classes and

        only module-level classes can be pickled by the default path.

        """

        if not self._deferred:

            return super(Model, self).__reduce__()

        data = self.__dict__

        if not self._deferred:

            class_id = self._meta.app_label, self._meta.object_name

            return model_unpickle, (class_id, [], simple_class_factory), data

        defers = []

        for field in self._meta.fields:

            if isinstance(self.__class__.__dict__.get(field.attname),

                          DeferredAttribute):

                defers.append(field.attname)

        model = self._meta.proxy_for_model

        return (model_unpickle, (model, defers), data)

        class_id = model._meta.app_label, model._meta.object_name

        return (model_unpickle, (class_id, defers, deferred_class_factory), data)



    def _get_pk_val(self, meta=None):

        if not meta:
@@ -1008,12 +1010,22 @@ def get_absolute_url(opts, func, self, *args, **kwargs): 
class Empty(object):

    pass



def simple_class_factory(model, attrs):

    """

    Needed for dynamic classes.

    """

    return model



def model_unpickle(model, attrs):

def model_unpickle(model_id, attrs, factory):

    """

    Used to unpickle Model subclasses with deferred fields.

    """

    cls = deferred_class_factory(model, attrs)

    if isinstance(model_id, tuple):

        model = get_model(*model_id)

    else:

        # Backwards compat - the model was cached directly in earlier versions.

        model = model_id

    cls = factory(model, attrs)

    return cls.__new__(cls)

model_unpickle.__safe_for_unpickle__ = True

tests/queryset_pickle/models.py

+10 −0
Original line number Diff line number Diff line
@@ -36,3 +36,13 @@ class Happening(models.Model): 
    number2 = models.IntegerField(blank=True, default=Numbers.get_static_number)

    number3 = models.IntegerField(blank=True, default=Numbers.get_class_number)

    number4 = models.IntegerField(blank=True, default=nn.get_member_number)



class Container(object):

    # To test pickling we need a class that isn't defined on module, but

    # is still available from app-cache. So, the Container class moves

    # SomeModel outside of module level

    class SomeModel(models.Model):

        somefield = models.IntegerField()



class M2MModel(models.Model):

    groups = models.ManyToManyField(Group)

tests/queryset_pickle/tests.py

+42 −1
Original line number Diff line number Diff line
@@ -3,9 +3,10 @@ from __future__ import absolute_import 
import pickle

import datetime



from django.db import models

from django.test import TestCase



from .models import Group, Event, Happening

from .models import Group, Event, Happening, Container, M2MModel





class PickleabilityTestCase(TestCase):
@@ -49,3 +50,43 @@ class PickleabilityTestCase(TestCase): 
        # can't just use assertEqual(original, unpickled)

        self.assertEqual(original.__class__, unpickled.__class__)

        self.assertEqual(original.args, unpickled.args)



    def test_model_pickle(self):

        """

        Test that a model not defined on module level is pickleable.

        """

        original = Container.SomeModel(pk=1)

        dumped = pickle.dumps(original)

        reloaded = pickle.loads(dumped)

        self.assertEqual(original, reloaded)

        # Also, deferred dynamic model works

        Container.SomeModel.objects.create(somefield=1)

        original = Container.SomeModel.objects.defer('somefield')[0]

        dumped = pickle.dumps(original)

        reloaded = pickle.loads(dumped)

        self.assertEqual(original, reloaded)

        self.assertEqual(original.somefield, reloaded.somefield)



    def test_model_pickle_m2m(self):

        """

        Test intentionally the automatically created through model.

        """

        m1 = M2MModel.objects.create()

        g1 = Group.objects.create(name='foof')

        m1.groups.add(g1)

        m2m_through = M2MModel._meta.get_field_by_name('groups')[0].rel.through

        original = m2m_through.objects.get()

        dumped = pickle.dumps(original)

        reloaded = pickle.loads(dumped)

        self.assertEqual(original, reloaded)



    def test_model_pickle_dynamic(self):

        class Meta:

            proxy = True

        dynclass = type("DynamicEventSubclass", (Event, ),

                        {'Meta': Meta, '__module__': Event.__module__})

        original = dynclass(pk=1)

        dumped = pickle.dumps(original)

        reloaded = pickle.loads(dumped)

        self.assertEqual(original, reloaded)

        self.assertIs(reloaded.__class__, dynclass)