瀏覽代碼

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

Thanks Tim Graham for the review.
Andrei Fokau 9 年之前
父節點
當前提交
998894e1b9

+ 14 - 0
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)

+ 2 - 0
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

+ 1 - 1
docs/releases/1.10.txt

@@ -174,7 +174,7 @@ Management Commands
 Migrations
 ^^^^^^^^^^
 
-* ...
+* Added support for serialization of ``enum.Enum`` objects.
 
 Models
 ^^^^^^

+ 5 - 0
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)

+ 60 - 0
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)

+ 1 - 0
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