From 56ec31d33362b8c297a10e6b0cef9577ab84298c Mon Sep 17 00:00:00 2001
From: Mikhail Podgurskiy <kmmbvnr@gmail.com>
Date: Sat, 14 May 2016 11:39:14 +0600
Subject: [PATCH] Support for Non-deterministic FSM? #129 and callble target
 state #61

---
 django_fsm/__init__.py                        | 41 +++++++++++++++++++
 tests/testapp/tests/test_multi_resultstate.py | 40 ++++++++++++++++++
 2 files changed, 81 insertions(+)
 create mode 100644 tests/testapp/tests/test_multi_resultstate.py

diff --git a/django_fsm/__init__.py b/django_fsm/__init__.py
index e62d83ef..17433545 100644
--- a/django_fsm/__init__.py
+++ b/django_fsm/__init__.py
@@ -65,6 +65,10 @@ def __init__(self, *args, **kwargs):
         super(TransitionNotAllowed, self).__init__(*args, **kwargs)
 
 
+class InvalidResultState(Exception):
+    """Raised when we got invalid result state"""
+
+
 class ConcurrentTransition(Exception):
     """
     Raised when the transition cannot be executed because the
@@ -318,6 +322,10 @@ def change_state(self, instance, method, *args, **kwargs):
         try:
             result = method(instance, *args, **kwargs)
             if next_state is not None:
+                if hasattr(next_state, 'get_state'):
+                    next_state = next_state.get_state(
+                        instance, transition, result,
+                        args=args, kwargs=kwargs)
                 self.set_proxy(instance, next_state)
                 self.set_state(instance, next_state)
         except Exception as exc:
@@ -547,3 +555,36 @@ def has_transition_perm(bound_method, user):
     return (meta.has_transition(current_state) and
             meta.conditions_met(im_self, current_state) and
             meta.has_transition_perm(im_self, current_state, user))
+
+
+class State(object):
+    def get_state(self, model, transition, result, args=[], kwargs={}):
+        raise NotImplementedError
+
+
+class RETURN_VALUE(State):
+    def __init__(self, *allowed_states):
+        self.allowed_states = allowed_states if allowed_states else None
+
+    def get_state(self, model, transition, result, args=[], kwargs={}):
+        if self.allowed_states is not None:
+            if result not in self.allowed_states:
+                raise InvalidResultState(
+                    '{} is not in list of allowed states\n{}'.format(
+                        result, self.allowed_states))
+        return result
+
+
+class GET_STATE(State):
+    def __init__(self, func, states=None):
+        self.func = func
+        self.allowed_states = states
+
+    def get_state(self, model, transition, result, args=[], kwargs={}):
+        result_state = self.func(model, *args, **kwargs)
+        if self.allowed_states is not None:
+            if result_state not in self.allowed_states:
+                raise InvalidResultState(
+                    '{} is not in list of allowed states\n{}'.format(
+                        result, self.allowed_states))
+        return result_state
diff --git a/tests/testapp/tests/test_multi_resultstate.py b/tests/testapp/tests/test_multi_resultstate.py
new file mode 100644
index 00000000..e7b385df
--- /dev/null
+++ b/tests/testapp/tests/test_multi_resultstate.py
@@ -0,0 +1,40 @@
+from django.db import models
+from django.test import TestCase
+from django_fsm import FSMField, transition, RETURN_VALUE, GET_STATE
+
+
+class MultiResultTest(models.Model):
+    state = FSMField(default='new')
+
+    @transition(
+        field=state,
+        source='new',
+        target=RETURN_VALUE('for_moderators', 'published'))
+    def publish(self, is_public=False):
+        return 'published' if is_public else 'for_moderators'
+
+    @transition(
+        field=state,
+        source='for_moderators',
+        target=GET_STATE(
+            lambda self, allowed: 'published' if allowed else 'rejected',
+            states=['published', 'rejected']
+        )
+    )
+    def moderate(self, allowed):
+        pass
+
+    class Meta:
+        app_label = 'testapp'
+
+
+class Test(TestCase):
+    def test_return_state_succeed(self):
+        instance = MultiResultTest()
+        instance.publish(is_public=True)
+        self.assertEqual(instance.state, 'published')
+
+    def test_get_state_succeed(self):
+        instance = MultiResultTest(state='for_moderators')
+        instance.moderate(allowed=False)
+        self.assertEqual(instance.state, 'rejected')