Add class substitution for state #41

django_fsm/__init__.py

@@ -6,6 +6,7 @@
 from functools import wraps
 
 from django.db import models
+from django.db.models.loading import get_model
 from django.db.models.signals import class_prepared
 from django.utils.functional import curry
 from django_fsm.signals import pre_transition, post_transition
@@ -164,6 +165,9 @@ def __get__(self, instance, type=None):
     def __set__(self, instance, value):
         if self.field.protected and self.field.name in instance.__dict__:
             raise AttributeError('Direct {} modification is not allowed'.format(self.field.name))
+
+        # Update state
+        self.field.set_proxy(instance, value)
         self.field.set_state(instance, value)
 
 
@@ -173,6 +177,19 @@ class FSMFieldMixin(object):
     def __init__(self, *args, **kwargs):
         self.protected = kwargs.pop('protected', False)
         self.transitions = {}  # cls -> (transitions name -> method)
+        self.state_proxy = {}  # state -> ProxyClsRef
+
+        state_choices = kwargs.pop('state_choices', None)
+        choices = kwargs.get('choices', None)
+        if state_choices is not None and choices is not None:
+            raise ValueError('Use one of choices or state_choces value')
+
+        if state_choices is not None:
+            choices = []
+            for state, title, proxy_cls_ref in state_choices:
+                choices.append((state, title))
+                self.state_proxy[state] = proxy_cls_ref
+            kwargs['choices'] = choices
 
         super(FSMFieldMixin, self).__init__(*args, **kwargs)
 
@@ -188,6 +205,26 @@ def get_state(self, instance):
     def set_state(self, instance, state):
         instance.__dict__[self.name] = state
 
+    def set_proxy(self, instance, state):
+        """
+        Change class
+        """
+        if state in self.state_proxy:
+            state_proxy = self.state_proxy[state]
+
+            try:
+                app_label, model_name = state_proxy.split(".")
+            except ValueError:
+                # If we can't split, assume a model in current app
+                app_label = instance._meta.app_label
+                model_name = state_proxy
+
+            model = get_model(app_label, model_name)
+            if model is None:
+                raise ValueError('No model found {}'.format(state_proxy))
+
+            instance.__class__ = model
+
     def change_state(self, instance, method, *args, **kwargs):
         meta = method._django_fsm
         method_name = method.__name__
@@ -211,6 +248,7 @@ def change_state(self, instance, method, *args, **kwargs):
 
         result = method(instance, *args, **kwargs)
         if next_state:
+            self.set_proxy(instance, next_state)
             self.set_state(instance, next_state)
 
         post_transition.send(**signal_kwargs)

django_fsm/tests/test_basic_transitions.py

@@ -120,6 +120,7 @@ def test_available_conditions(self):
 
     def test_all_conditions(self):
         transitions = self.model.get_all_state_transitions()
+
         actual = set((transition.source, transition.target) for transition in transitions)
         expected = set([('*', 'moderated'),
                         ('new', 'published'),

tests/settings.py

@@ -3,10 +3,9 @@
 DATABASE_ENGINE = 'sqlite3'
 SECRET_KEY = 'nokey'
 
-DATABASE_ENGINE = 'sqlite3'
 DATABASES = {
     'default': {
-        'ENGINE': 'django.db.backends.%s' % DATABASE_ENGINE,
+        'ENGINE': 'django.db.backends.sqlite3',
         }
 }
 
@@ -15,11 +14,3 @@
     'django_jenkins.tasks.run_pep8',
     'django_jenkins.tasks.run_pyflakes'
 )
-
-if __name__ == "__main__":
-    import sys, test_runner as settings
-    from django.core.management import execute_manager
-
-    if len(sys.argv) == 1:
-            sys.argv += ['test'] + list(PROJECT_APPS)
-    execute_manager(settings)

tests/testapp/tests/test_state_transitions.py

@@ -0,0 +1,67 @@
+from django.db import models
+from django.test import TestCase
+from django_fsm import FSMField, transition
+
+
+class Insect(models.Model):
+    class STATE:
+        CATERPILLAR = 'CTR'
+        BUTTERFLY = 'BTF'
+
+    STATE_CHOICES = ((STATE.CATERPILLAR, 'Caterpillar', 'Caterpillar'),
+                     (STATE.BUTTERFLY, 'Butterfly', 'Butterfly'))
+
+    state = FSMField(default=STATE.CATERPILLAR, state_choices=STATE_CHOICES)
+
+    @transition(field=state, source=STATE.CATERPILLAR, target=STATE.BUTTERFLY)
+    def cocoon(self):
+        pass
+
+    def fly(self):
+        raise NotImplementedError
+
+    def crawl(self):
+        raise NotImplementedError
+
+    class Meta:
+        app_label = 'testapp'
+
+
+class Caterpillar(Insect):
+    def crawl(self):
+        """
+        Do crawl
+        """
+
+    class Meta:
+        app_label = 'testapp'
+        proxy = True
+
+
+class Butterfly(Insect):
+    def fly(self):
+        """
+        Do fly
+        """
+
+    class Meta:
+        app_label = 'testapp'
+        proxy = True
+
+
+class TestStateProxy(TestCase):
+    def test_initial_proxy_set_succeed(self):
+        insect = Insect()
+        self.assertTrue(isinstance(insect, Caterpillar))
+
+    def test_transition_proxy_set_succeed(self):
+        insect = Insect()
+        insect.cocoon()
+        self.assertTrue(isinstance(insect, Butterfly))
+
+    def test_load_proxy_set(self):
+        Insect.objects.create(state=Insect.STATE.CATERPILLAR)
+        Insect.objects.create(state=Insect.STATE.BUTTERFLY)
+
+        insects = Insect.objects.all()
+        self.assertEqual(set([Caterpillar, Butterfly]), set(insect.__class__ for insect in insects))