Commit 998894e1 authored by Andrei Fokau's avatar Andrei Fokau Committed by Tim Graham
Browse files

Fixed #25764 -- Added support for serialization of enum.Enum in migrations. 

Thanks Tim Graham for the review.
parent 65764a93

django/db/migrations/writer.py

+14 −0
Original line number Diff line number Diff line
@@ -25,6 +25,12 @@ from django.utils.module_loading import module_dir 
from django.utils.timezone import now, utc

from django.utils.version import get_docs_version



try:

    import enum

except ImportError:

    # No support on Python 2 if enum34 isn't installed.

    enum = None





class SettingsReference(str):

    """
@@ -376,6 +382,14 @@ class MigrationWriter(object): 
                imports.update(v_imports)

                strings.append((k_string, v_string))

            return "{%s}" % (", ".join("%s: %s" % (k, v) for k, v in strings)), imports

        # Enums

        elif enum and isinstance(value, enum.Enum):

            enum_class = value.__class__

            module = enum_class.__module__

            imports = {"import %s" % module}

            v_string, v_imports = cls.serialize(value.value)

            imports.update(v_imports)

            return "%s.%s(%s)" % (module, enum_class.__name__, v_string), imports

        # Datetimes

        elif isinstance(value, datetime.datetime):

            value_repr = cls.serialize_datetime(value)

docs/internals/contributing/writing-code/unit-tests.txt

+2 −0
Original line number Diff line number Diff line
@@ -138,6 +138,7 @@ dependencies: 


*  bcrypt_

*  docutils_

*  enum34_ (Python 2 only)

*  geoip2_

*  jinja2_ 2.7+

*  numpy_
@@ -171,6 +172,7 @@ associated tests will be skipped. 


.. _bcrypt: https://pypi.python.org/pypi/bcrypt

.. _docutils: https://pypi.python.org/pypi/docutils

.. _enum34: https://pypi.python.org/pypi/enum34

.. _geoip2: https://pypi.python.org/pypi/geoip2

.. _jinja2: https://pypi.python.org/pypi/jinja2

.. _numpy: https://pypi.python.org/pypi/numpy

docs/releases/1.10.txt

+1 −1
Original line number Diff line number Diff line
@@ -174,7 +174,7 @@ Management Commands 
Migrations

^^^^^^^^^^



* ...

* Added support for serialization of ``enum.Enum`` objects.



Models

^^^^^^

docs/topics/migrations.txt

+5 −0
Original line number Diff line number Diff line
@@ -645,6 +645,7 @@ Django can serialize the following: 
- ``datetime.date``, ``datetime.time``, and ``datetime.datetime`` instances

  (include those that are timezone-aware)

- ``decimal.Decimal`` instances

- ``enum.Enum`` instances

- ``functools.partial`` instances which have serializable ``func``, ``args``,

  and ``keywords`` values.

- Any Django field
@@ -656,6 +657,10 @@ Django can serialize the following: 


    Serialization support for `functools.partial` was added.



.. versionchanged:: 1.10



    Serialization support for ``enum.Enum`` was added.



Django can serialize the following on Python 3 only:



- Unbound methods used from within the class body (see below)

tests/migrations/test_writer.py

+60 −0
Original line number Diff line number Diff line
@@ -28,6 +28,11 @@ from django.utils.translation import ugettext_lazy as _ 


from .models import FoodManager, FoodQuerySet



try:

    import enum

except ImportError:

    enum = None





class TestModel1(object):

    def upload_to(self):
@@ -229,6 +234,61 @@ class WriterTests(SimpleTestCase): 
            ("[list, tuple, dict, set, frozenset]", set())

        )



    @unittest.skipUnless(enum, "enum34 is required on Python 2")

    def test_serialize_enums(self):

        class TextEnum(enum.Enum):

            A = 'a-value'

            B = 'value-b'



        class BinaryEnum(enum.Enum):

            A = b'a-value'

            B = b'value-b'



        class IntEnum(enum.IntEnum):

            A = 1

            B = 2



        self.assertSerializedResultEqual(

            TextEnum.A,

            ("migrations.test_writer.TextEnum('a-value')", {'import migrations.test_writer'})

        )

        self.assertSerializedResultEqual(

            BinaryEnum.A,

            ("migrations.test_writer.BinaryEnum(b'a-value')", {'import migrations.test_writer'})

        )

        self.assertSerializedResultEqual(

            IntEnum.B,

            ("migrations.test_writer.IntEnum(2)", {'import migrations.test_writer'})

        )



        field = models.CharField(default=TextEnum.B, choices=[(m.value, m) for m in TextEnum])

        string = MigrationWriter.serialize(field)[0]

        self.assertEqual(

            string,

            "models.CharField(choices=["

            "('a-value', migrations.test_writer.TextEnum('a-value')), "

            "('value-b', migrations.test_writer.TextEnum('value-b'))], "

            "default=migrations.test_writer.TextEnum('value-b'))"

        )

        field = models.CharField(default=BinaryEnum.B, choices=[(m.value, m) for m in BinaryEnum])

        string = MigrationWriter.serialize(field)[0]

        self.assertEqual(

            string,

            "models.CharField(choices=["

            "(b'a-value', migrations.test_writer.BinaryEnum(b'a-value')), "

            "(b'value-b', migrations.test_writer.BinaryEnum(b'value-b'))], "

            "default=migrations.test_writer.BinaryEnum(b'value-b'))"

        )

        field = models.IntegerField(default=IntEnum.A, choices=[(m.value, m) for m in IntEnum])

        string = MigrationWriter.serialize(field)[0]

        self.assertEqual(

            string,

            "models.IntegerField(choices=["

            "(1, migrations.test_writer.IntEnum(1)), "

            "(2, migrations.test_writer.IntEnum(2))], "

            "default=migrations.test_writer.IntEnum(1))"

        )



    def test_serialize_functions(self):

        with six.assertRaisesRegex(self, ValueError, 'Cannot serialize function: lambda'):

            self.assertSerializedEqual(lambda x: 42)