optional choices

Author: marcinn

dsm/__init__.py

@@ -1,4 +1,3 @@
-
 import collections
 import observable
 
@@ -43,8 +42,9 @@ def has_state(self, state):
     def register(self, from_state, value, to_state):
         if from_state in self._states and value in self._states[from_state]:
             raise AlreadyRegistered(
-                'Transition for `%s` is already registered for state `%s`' % (
-                    value, from_state))
+                "Transition for `%s` is already registered for state `%s`"
+                % (value, from_state)
+            )
         self._states[from_state][value] = to_state
         self._allstates.update([from_state, to_state])
 
@@ -60,16 +60,16 @@ def register_fallback(self, from_state, to_state):
 
         if from_state in self._fallbacks:
             raise AlreadyRegistered(
-                'Fallback transition for `%s` '
-                'is already registered' % from_state)
+                "Fallback transition for `%s` " "is already registered" % from_state
+            )
 
         self._fallbacks[from_state] = to_state
         self._allstates.update([from_state, to_state])
 
     def can(self, value, current_state):
         return bool(
-                self._states.get(current_state) and
-                self._states[current_state].get(value))
+            self._states.get(current_state) and self._states[current_state].get(value)
+        )
 
     def execute(self, value, current_state):
         try:
@@ -79,8 +79,9 @@ def execute(self, value, current_state):
                 return self._fallbacks[current_state]
             except KeyError:
                 raise UnknownTransition(
-                    'Can not find transition for `%s` in state `%s`' % (
-                                                    value, current_state))
+                    "Can not find transition for `%s` in state `%s`"
+                    % (value, current_state)
+                )
 
 
 class MetaMachine(type):
@@ -93,36 +94,38 @@ def __new__(cls, name, bases, attrs):
             cls.add_exception_classes(new_class)
             return new_class
 
-        meta = attrs.pop('Meta', None)
+        meta = attrs.pop("Meta", None)
 
         class Options:
             def __init__(self, meta):
                 self.transitions = Transitions(
-                        transitions=getattr(meta, 'transitions', None),
-                        fallbacks=getattr(meta, 'fallbacks', None))
-                self.initial = getattr(meta, 'initial', None)
+                    transitions=getattr(meta, "transitions", None),
+                    fallbacks=getattr(meta, "fallbacks", None),
+                )
+                self.initial = getattr(meta, "initial", None)
 
         new_class = super_new(cls, name, bases, {})
         cls.add_exception_classes(new_class)
-        setattr(new_class, '_meta', Options(meta))
+        setattr(new_class, "_meta", Options(meta))
 
         return new_class
 
     def add_exception_classes(new_class):
-        setattr(new_class, 'FSMException', FSMException)
-        setattr(new_class, 'UnknownTransition', UnknownTransition)
+        setattr(new_class, "FSMException", FSMException)
+        setattr(new_class, "UnknownTransition", UnknownTransition)
 
 
 class StateMachine(metaclass=MetaMachine):
     def __init__(self, initial=None, transitions=None):
-        meta = getattr(self, '_meta', None)
+        meta = getattr(self, "_meta", None)
         self._eventhandler = observable.Observable()
-        self._transitions = transitions or getattr(
-                                meta, 'transitions', None) or Transitions()
-        self._initial = initial or getattr(meta, 'initial', None)
+        self._transitions = (
+            transitions or getattr(meta, "transitions", None) or Transitions()
+        )
+        self._initial = initial or getattr(meta, "initial", None)
         self._state = None
         self._inputhandlers = collections.defaultdict(list)
-        self._eventhandler.on('input', self._inputhandler)
+        self._eventhandler.on("input", self._inputhandler)
         self.reset()
 
     @property
@@ -133,11 +136,10 @@ def process(self, value):
         new_state = self._transitions.execute(value, self.state)
 
         if not self.state == new_state:
-            self._eventhandler.trigger(
-                    'change', state=new_state, previous=self.state)
+            self._eventhandler.trigger("change", state=new_state, previous=self.state)
 
         self._state = new_state
-        self._eventhandler.trigger('input', state=new_state, value=value)
+        self._eventhandler.trigger("input", state=new_state, value=value)
 
         return self.state
 
@@ -155,9 +157,8 @@ def reset(self):
 
         old_state = self._state
         self._state = self._initial
-        self._eventhandler.trigger(
-                'change', state=self._state, previous=old_state)
-        self._eventhandler.trigger('reset')
+        self._eventhandler.trigger("change", state=self._state, previous=old_state)
+        self._eventhandler.trigger("reset")
         return self.state
 
     def when(self, state, func):

dsm/fields.py

@@ -1,8 +1,9 @@
+from django.core.checks import Error
 from django.db import models
 
 from . import StateMachine, Transitions
 
-__all__ = ['StateMachineField']
+__all__ = ["StateMachineField"]
 
 
 class MachineState:
@@ -49,13 +50,11 @@ def __get__(self, instance, cls=None):
             return value
         if isinstance(value, MachineState):
             return value
-        return MachineState(
-            instance, self.field, self.field._create_fsm(value))
+        return MachineState(instance, self.field, self.field._create_fsm(value))
 
     def __set__(self, instance, value):
         if value is not None:
-            value = MachineState(
-                instance, self.field, self.field._create_fsm(value))
+            value = MachineState(instance, self.field, self.field._create_fsm(value))
         instance.__dict__[self.field.name] = value
         """
         current_value = instance.__dict__.get(self.field.name)
@@ -74,12 +73,18 @@ class StateMachineFieldMixin:
     descriptor_class = StateDescriptor
 
     def __init__(self, transitions, *args, **kwargs):
-        self.transitions = transitions
+        if isinstance(transitions, Transitions):
+            self.transitions = transitions
+        else:
+            self.transitions = Transitions(transitions)
+
+        if "choices" not in kwargs:
+            kwargs["choices"] = [(s, s) for s in self.transitions._allstates]
+
         super().__init__(*args, **kwargs)
 
     def _create_fsm(self, initial):
-        return StateMachine(
-            initial=initial, transitions=Transitions(self.transitions))
+        return StateMachine(initial=initial, transitions=self.transitions)
 
     def get_prep_value(self, value):
         if value is None:
@@ -92,8 +97,22 @@ def contribute_to_class(self, cls, name, **kwargs):
 
 
 class StateMachineField(StateMachineFieldMixin, models.CharField):
+    def _check_choices(self, **kwargs):
+        all_states = self.transitions._allstates
+        choice_states = {x[0] for x in self.choices}
+        if not all_states.issubset(choice_states):
+            return [
+                Error(
+                    "Following states are not defined in choices: %s"
+                    % (", ".join(all_states - choice_states)),
+                    obj=self,
+                    id="dsm.E001",
+                )
+            ]
+        return []
+
     def deconstruct(self):
         name, path, args, kwargs = super().deconstruct()
-        path = 'dsm.fields.StateMachineField'
+        path = "dsm.fields.StateMachineField"
         args.insert(0, self.transitions)
         return name, path, args, kwargs

requirements-dev.txt

@@ -4,3 +4,5 @@ ipdb
 wheel
 pytest
 pytest-django
+isort
+black

tests/models.py

@@ -1,14 +1,14 @@
 from django.db import models
-from django.utils.translation import gettext_lazy as _
 from dsm.fields import StateMachineField
 
+
 class Order(models.Model):
     class Status(models.TextChoices):
-        NEW = 'new', _('New')
-        PROCESSING = 'processing', _('Processing')
-        SENDING = 'sending', _('Sending')
-        FINISHED = 'finished', _('Finished')
-        CANCELLED = 'cancelled', _('Cancelled')
+        NEW = 'new', 'New'
+        PROCESSING = 'processing', 'Processing'
+        SENDING = 'sending', 'Sending'
+        FINISHED = 'finished', 'Finished'
+        CANCELLED = 'cancelled', 'Cancelled'
 
     status = StateMachineField(
         transitions=(
@@ -25,3 +25,27 @@ class Status(models.TextChoices):
 
     class Meta:
         app_label = 'tests'
+
+
+class OrderNoChoices(models.Model):
+    class Status(models.TextChoices):
+        NEW = 'new', 'New'
+        PROCESSING = 'processing', 'Processing'
+        SENDING = 'sending', 'Sending'
+        FINISHED = 'finished', 'Finished'
+        CANCELLED = 'cancelled', 'Cancelled'
+
+    status = StateMachineField(
+        transitions=(
+            (Status.NEW, ['confirm'], Status.PROCESSING),
+            (Status.PROCESSING, ['cancel'], Status.CANCELLED),
+            (Status.PROCESSING, ['send'], Status.SENDING),
+            (Status.SENDING, ['deliver'], Status.FINISHED),
+        ),
+        max_length=16,
+        db_index=True,
+        default=Status.NEW,
+    )
+
+    class Meta:
+        app_label = 'tests'