diff --git a/django/db/migrations/writer.py b/django/db/migrations/writer.py index aa9f393e4b..692beef384 100644 --- a/django/db/migrations/writer.py +++ b/django/db/migrations/writer.py @@ -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) diff --git a/docs/internals/contributing/writing-code/unit-tests.txt b/docs/internals/contributing/writing-code/unit-tests.txt index cdf7960ac5..a2013d03bb 100644 --- a/docs/internals/contributing/writing-code/unit-tests.txt +++ b/docs/internals/contributing/writing-code/unit-tests.txt @@ -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 diff --git a/docs/releases/1.10.txt b/docs/releases/1.10.txt index e4cae5a983..9eb6e1d93a 100644 --- a/docs/releases/1.10.txt +++ b/docs/releases/1.10.txt @@ -174,7 +174,7 @@ Management Commands Migrations ^^^^^^^^^^ -* ... +* Added support for serialization of ``enum.Enum`` objects. Models ^^^^^^ diff --git a/docs/topics/migrations.txt b/docs/topics/migrations.txt index 546c1a9807..44b3afe5f0 100644 --- a/docs/topics/migrations.txt +++ b/docs/topics/migrations.txt @@ -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) diff --git a/tests/migrations/test_writer.py b/tests/migrations/test_writer.py index 12df772df5..e6928b6bbb 100644 --- a/tests/migrations/test_writer.py +++ b/tests/migrations/test_writer.py @@ -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) diff --git a/tests/requirements/py2.txt b/tests/requirements/py2.txt index 3bfc36436d..0ed845bb7e 100644 --- a/tests/requirements/py2.txt +++ b/tests/requirements/py2.txt @@ -1,4 +1,5 @@ -r base.txt +enum34 # Due to https://github.com/linsomniac/python-memcached/issues/79 in newer versions. python-memcached <= 1.53 mock