Source

django-statemachine / statemachine / fields.py

Full commit
from copy import deepcopy
from django.db import models
from fsm import State, FSM
from django.forms.forms import pretty_name

DEFAULT_STATE = "start"

class FSM_StateField(models.Field):
    """
    Django models.Field subclass used to store a representation of
    state from fsm.FSM

    The underlying database value is a char(50)

    Usage::

        class Foo(models.Model):
            blah = FSM_StateField(fsmFSM_class)

        foo = Foo.objects.create()
        print foo.blah.state
        foo.blah.change('new_state')
        foo.save()
    """

    __metaclass__ = models.SubfieldBase
    description = "Finite State Machine Field"

    def __init__(self, machine, default_choices_all=False, *args, **kwargs):
        """
        machine: instance of fsm.FSM
        default_choices_all: 

        * If True sets choices to the results of all_state_choices
        * If False set choices to the results of exit_state_choices

        This constructor stores a copy of the class definition in
        __machine_class and instantiates it in self.machine (*setting
        an initial state from the database if available*) when
        returning a python object.
        """

        self.__machine_class = machine
        self.default_choices_all = default_choices_all
        defaults = {'max_length': 50, 'default' : DEFAULT_STATE}
        defaults.update(kwargs)
        super(FSM_StateField, self).__init__(self, **defaults)

    def __unicode__(self):
        return self.name

    def db_type(self, connection):
        return "char(50)"

    def to_python(self, value):
        self.machine = self.__machine_class()
        try:
            name = value.state
        except AttributeError:
            name = value

        self.machine.set_initial_state(name)
        if self.default_choices_all:
            self._choices = self.all_state_choices()
        else:
            self._choices = self.exit_state_choices()
        return self.machine

    def validate(self, value, model_instance):
        pass

    def get_prep_value(self, value):
        return value

    def get_db_prep_value(self, value, connection, prepared=False):
        try:
            name = value.state
        except AttributeError:
            name = value
        return name

    def all_state_choices(self):
        """
        Returns a choices list of all the available states, useful for the admin interface.
        """
        return [(name, pretty_name(name)) for name, state in self.machine.states.items()]

    def exit_state_choices(self):
        """
        Returns a choices list of only the available exit states from the current state.
        """
        state_choices = []
        if self.machine.state:
            state_choices.append((self.machine.state,self.machine.state))
            state_choices.extend([(name,name) for name in self.machine.states[self.machine.state].exit_states])

        return state_choices