Merge pull request piccolo-orm#114 from piccolo-orm/choice_defaults

Choice defaults

Author: dantownsend

piccolo/apps/migrations/auto/serialisation.py

@@ -73,6 +73,20 @@ def __repr__(self):
         return f"{self.instance.__class__.__name__}({args})"
 
 
+@dataclass
+class SerialisedEnumInstance:
+    instance: Enum
+
+    def __hash__(self):
+        return hash(self.__repr__())
+
+    def __eq__(self, other):
+        return self.__hash__() == other.__hash__()
+
+    def __repr__(self):
+        return f"{self.instance.__class__.__name__}.{self.instance.name}"
+
+
 @dataclass
 class SerialisedTableType:
     table_type: t.Type[Table]
@@ -241,14 +255,29 @@ def serialise_params(params: t.Dict[str, t.Any]) -> SerialisedParams:
             extra_imports.append(Import(module="decimal", target="Decimal"))
             continue
 
-        # Enums
+        # Enum instances
         if isinstance(value, Enum):
-            # Enums already have a good __repr__.
-            extra_imports.append(
-                Import(
-                    module=value.__module__, target=value.__class__.__name__,
+            if value.__module__.startswith("piccolo"):
+                # It's an Enum defined within Piccolo, so we can safely import
+                # it.
+                params[key] = SerialisedEnumInstance(instance=value)
+                extra_imports.append(
+                    Import(
+                        module=value.__module__,
+                        target=value.__class__.__name__,
+                    )
                 )
-            )
+            else:
+                # It's a user defined Enum, so we'll insert the raw value.
+                enum_serialised_params: SerialisedParams = serialise_params(
+                    params={key: value.value}
+                )
+                params[key] = enum_serialised_params.params[key]
+                extra_imports.extend(enum_serialised_params.extra_imports)
+                extra_definitions.extend(
+                    enum_serialised_params.extra_definitions
+                )
+
             continue
 
         # Enum types
@@ -329,5 +358,7 @@ def deserialise_params(params: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
             params[key] = value.table_type
         elif isinstance(value, SerialisedEnumType):
             params[key] = value.enum_type
+        elif isinstance(value, SerialisedEnumInstance):
+            params[key] = value.instance
 
     return params

tests/apps/migrations/auto/test_serialisation.py

@@ -2,6 +2,7 @@
 from unittest import TestCase
 
 from piccolo.apps.migrations.auto.serialisation import serialise_params
+from piccolo.columns.base import OnDelete
 from piccolo.columns.choices import Choice
 from piccolo.columns.column_types import Varchar
 from piccolo.columns.defaults import DateNow, TimeNow, TimestampNow, UUID4
@@ -148,3 +149,36 @@ class Choices(Enum):
                 "from enum import Enum",
             },
         )
+
+    def test_custom_enum_instance(self):
+        """
+        Make sure custom Enum instances can be serialised properly. An example
+        is when a user defines a choices Enum, and then sets the default to
+        one of those choices.
+        """
+
+        class Choices(Enum):
+            a = 1
+            b = 2
+
+        serialised = serialise_params(params={"default": Choices.a})
+
+        self.assertEqual(serialised.params["default"], 1)
+        self.assertEqual(serialised.extra_imports, [])
+        self.assertEqual(serialised.extra_definitions, [])
+
+    def test_builtin_enum_instance(self):
+        """
+        Make sure Enum instances defiend in Piccolo can be serialised properly
+        - for example, with on_delete.
+        """
+        serialised = serialise_params(params={"on_delete": OnDelete.cascade})
+
+        self.assertEqual(
+            serialised.params["on_delete"].__repr__(), "OnDelete.cascade"
+        )
+        self.assertEqual(
+            [i.__repr__() for i in serialised.extra_imports],
+            ["from piccolo.columns.base import OnDelete"],
+        )
+        self.assertEqual(serialised.extra_definitions, [])