Resolve "extend worker class"

Closes #183 (closed)

smash/web/forms/worker_form.py

+import logging
+
+from django import forms
+from django.forms import ModelForm
+
+from web.models import Worker, WorkerStudyRole
+from web.models.constants import GLOBAL_STUDY_ID
+from web.models.worker import role_choices_by_worker_type, worker_type_by_worker
+from web.models.worker_study_role import WORKER_STAFF
+
+logger = logging.getLogger(__name__)
+
+
+class WorkerForm(ModelForm):
+    class Meta:
+        model = Worker
+        exclude = ['appointments']
+
+    def __init__(self, *args, **kwargs):
+        worker_type = kwargs.pop('worker_type', WORKER_STAFF)
+        super(WorkerForm, self).__init__(*args, **kwargs)
+        instance = getattr(self, 'instance', None)
+
+        initial_role = None
+        if instance is not None:
+            worker_type = worker_type_by_worker(instance)
+            roles = WorkerStudyRole.objects.filter(worker=instance, study_id=GLOBAL_STUDY_ID)
+
+            if roles.count() > 0:
+                initial_role = roles[0].role
+
+        choices = role_choices_by_worker_type(worker_type)
+        self.fields['role'] = forms.ChoiceField(label='Role', choices=choices)
+        self.fields['role'].initial = initial_role
+
+        if worker_type == WORKER_STAFF:
+            del self.fields['voucher_types']
+
+    def save(self, commit=True):
+        instance = super(WorkerForm, self).save(commit)
+        roles = WorkerStudyRole.objects.filter(worker=instance, study_id=GLOBAL_STUDY_ID)
+        if roles.count() > 0:
+            roles.update(role=self.cleaned_data['role'])
+        else:
+            WorkerStudyRole.objects.create(worker=instance, study_id=GLOBAL_STUDY_ID, role=self.cleaned_data['role'])