Wizard: fix M2M data storage in formset

1 files changed, 130 insertions, 82 deletions

diff --git a/ishtar_common/wizards.py b/ishtar_common/wizards.py
index ca4937ec3..5f3da1130 100644
--- a/ishtar_common/wizards.py
+++ b/ishtar_common/wizards.py
@@ -231,11 +231,14 @@ class Wizard(NamedUrlWizardView):
                 self.ignore_init_steps]
 
     def get_context_data(self, form, **kwargs):
-        """Add previous, next and current steps to manage the wizard path"""
+        """
+        Add previous, next and current steps to manage the wizard path
+        """
         context = super(Wizard, self).get_context_data(form)
         self.request.session['CURRENT_ACTION'] = self.get_wizard_name()
         step = self.steps.first
         current_step = self.steps.current
+
         dct = {'current_step_label': self.form_list[current_step].form_label,
                'wizard_label': self.get_label(),
                'current_object': self.get_current_object(),
@@ -243,9 +246,11 @@ class Wizard(NamedUrlWizardView):
                    [k for k in self.main_item_select_keys
                     if current_step.startswith(k)]) if current_step else False
                }
+
         context.update(dct)
         if step == current_step:
             return context
+
         previous_steps, next_steps, previous_step_counter = [], [], 0
         while step:
             if step == current_step \
@@ -257,6 +262,7 @@ class Wizard(NamedUrlWizardView):
             if previous_step_counter >= len(self.steps):
                 break
             step = self.steps.all[previous_step_counter]
+
         context.update({'previous_steps': previous_steps,
                         'previous_step_counter': previous_step_counter})
         storage = self.storage
@@ -264,15 +270,20 @@ class Wizard(NamedUrlWizardView):
         # if self.modification or True:
         next_step = self.steps.first
         current_step_passed, no_next = False, False
+
         # force rechecking of conditions
         self.get_form_list()
         last_step_is_available = True
+
         while next_step:
+
             # check if the form is initialized otherwise initialize it
+            # put initialized data in session
             if self.modification and not storage.get_step_data(next_step):
                 values = self.get_form_initial(next_step)
                 prefixed_values = MultiValueDict()
                 if not isinstance(values, list):
+                    # simple form
                     for key in values:
                         form_key = next_step + '-' + key
                         if isinstance(values, MultiValueDict):
@@ -281,11 +292,17 @@ class Wizard(NamedUrlWizardView):
                         else:
                             prefixed_values[form_key] = values[key]
                 else:
+                    # formset
                     for formset_idx, v in enumerate(values):
                         prefix = u"-%d-" % formset_idx
                         for key in v:
                             form_key = next_step + prefix + key
-                            prefixed_values[form_key] = v[key]
+                            if isinstance(v, MultiValueDict):
+                                prefixed_values.setlist(form_key,
+                                                        v.getlist(key))
+                            else:
+                                prefixed_values[form_key] = v[key]
+
                 if not prefixed_values and \
                         next_step not in self.get_ignore_init_steps():
                     # simulate a non empty data for form that might be
@@ -293,6 +310,7 @@ class Wizard(NamedUrlWizardView):
                     prefixed_values['__non_empty_data'] = ''
                 self.prepare_serialization(prefixed_values)
                 storage.set_step_data(next_step, prefixed_values)
+
             if step == next_step:
                 current_step_passed = True
             elif current_step_passed:
@@ -1101,6 +1119,7 @@ class Wizard(NamedUrlWizardView):
         current_obj = self.get_current_object()
         current_step = self.steps.current
         request = self.request
+
         step_is_main_select = bool([k for k in self.main_item_select_keys
                                     if step.startswith(k)])
         if step_is_main_select and step in self.form_list \
@@ -1114,6 +1133,7 @@ class Wizard(NamedUrlWizardView):
                 return MultiValueDict({'pk': val})
         elif current_obj:
             return self.get_instanced_init(current_obj, step)
+
         current_form = self.form_list[current_step]
         if hasattr(current_form, 'currents'):
             initial = MultiValueDict()
@@ -1135,96 +1155,125 @@ class Wizard(NamedUrlWizardView):
             obj_name = prefixes[-2]
         return obj_name
 
+    def _get_instanced_init_for_form(self, obj, c_form):
+        """
+        Get initial data from an object: simple form
+        """
+        initial = MultiValueDict()
+        for base_field in c_form.base_fields.keys():
+            value = obj
+            base_model = None
+            if hasattr(c_form, 'base_model') and \
+                    base_field == c_form.base_model:
+                base_model = base_field
+            if hasattr(c_form, 'base_models') and \
+                    base_field in c_form.base_models:
+                base_model = base_field
+            if base_model:
+                key = base_model + 's'
+                initial.setlist(base_field, [
+                    unicode(val.pk) for val in getattr(value, key).all()])
+            else:
+                fields = base_field.split('__')
+                for field in fields:
+                    if callable(value):
+                        value = value()
+                    if not hasattr(value, field) or \
+                            getattr(value, field) is None:
+                        value = obj
+                        break
+                    value = getattr(value, field)
+            if hasattr(value, 'all') and callable(value.all):
+                if not value.count():
+                    continue
+                initial.setlist(base_field,
+                                [unicode(v.pk) for v in value.all()])
+                continue
+            if value == obj:
+                continue
+            if hasattr(value, 'pk'):
+                value = value.pk
+            if value in (True, False) \
+                    or isinstance(value, ImageFieldFile) \
+                    or isinstance(value, FileField):
+                initial[base_field] = value
+            elif value is not None:
+                initial[base_field] = unicode(value)
+        return initial
+
+    def _get_instanced_init_for_formset(self, obj, current_step, c_form):
+        """
+        Get initial data from an object: formset
+        """
+        initial = []
+        if hasattr(c_form.form, 'base_model'):
+            key = c_form.form.base_model + 's'
+        else:
+            key = current_step.split('-')[0]
+        if not hasattr(obj, key):
+            return initial
+        keys = c_form.form.base_fields.keys()
+        related = getattr(obj, key)
+        # manage through
+        through = False
+        if hasattr(related, 'through') and related.through:
+            related_set_name = str(
+                related.through.__name__ + '_set').lower()
+            if hasattr(obj, related_set_name):
+                through = True
+                related = getattr(obj, related_set_name)
+
+        query = related
+        if not through and not obj._meta.ordering:
+            query = query.order_by('pk')
+
+        for child_obj in query.all():
+            if not keys:
+                break
+            vals = MultiValueDict()
+            if len(keys) == 1:
+                # only one field: must be the id of the object
+                vals[keys[0]] = unicode(child_obj.pk)
+            else:
+                for field in keys:
+                    if hasattr(child_obj, field):
+                        value = getattr(child_obj, field)
+                        if hasattr(value, 'pk'):
+                            value = value.pk
+                        if value is not None:
+                            vals[field] = unicode(value)
+                    elif hasattr(child_obj, field + "s"):
+                        # M2M
+                        vals.setlist(field, [
+                            unicode(v.pk)
+                            for v in getattr(child_obj, field + "s").all()
+                        ])
+            if vals:
+                initial.append(vals)
+        return initial
+
     def get_instanced_init(self, obj, step=None):
-        """Get initial data from an init"""
+        """
+        Get initial data from an object
+        """
         current_step = step or self.steps.current
         c_form = self.form_list[current_step]
+
         # make the current object the default item for the session
         self.request.session[self.get_object_name(obj)] = unicode(obj.pk)
+
         initial = MultiValueDict()
+        # posted data or already in session
         if self.request.POST or \
                 (step in self.request.session[self.storage.prefix] and
                  self.request.session[self.storage.prefix]['step_data'][step]):
             return initial
+
         if hasattr(c_form, 'base_fields'):
-            for base_field in c_form.base_fields.keys():
-                value = obj
-                base_model = None
-                if hasattr(c_form, 'base_model') and \
-                   base_field == c_form.base_model:
-                    base_model = base_field
-                if hasattr(c_form, 'base_models') and \
-                   base_field in c_form.base_models:
-                    base_model = base_field
-                if base_model:
-                    key = base_model + 's'
-                    initial.setlist(base_field, [
-                        unicode(val.pk) for val in getattr(value, key).all()])
-                else:
-                    fields = base_field.split('__')
-                    for field in fields:
-                        if callable(value):
-                            value = value()
-                        if not hasattr(value, field) or \
-                                getattr(value, field) is None:
-                            value = obj
-                            break
-                        value = getattr(value, field)
-                if hasattr(value, 'all') and callable(value.all):
-                    if not value.count():
-                        continue
-                    initial.setlist(base_field,
-                                    [unicode(v.pk) for v in value.all()])
-                    continue
-                if value == obj:
-                    continue
-                if hasattr(value, 'pk'):
-                    value = value.pk
-                if value in (True, False) \
-                        or isinstance(value, ImageFieldFile) \
-                        or isinstance(value, FileField):
-                    initial[base_field] = value
-                elif value is not None:
-                    initial[base_field] = unicode(value)
+            return self._get_instanced_init_for_form(obj, c_form)
         elif hasattr(c_form, 'management_form'):
-            initial = []
-            if hasattr(c_form.form, 'base_model'):
-                key = c_form.form.base_model + 's'
-            else:
-                key = current_step.split('-')[0]
-            if not hasattr(obj, key):
-                return initial
-            keys = c_form.form.base_fields.keys()
-            related = getattr(obj, key)
-            # manage through
-            through = False
-            if hasattr(related, 'through') and related.through:
-                related_set_name = str(
-                    related.through.__name__ + '_set').lower()
-                if hasattr(obj, related_set_name):
-                    through = True
-                    related = getattr(obj, related_set_name)
-
-            query = related
-            if not through and not obj._meta.ordering:
-                query = query.order_by('pk')
-            for child_obj in query.all():
-                if not keys:
-                    break
-                vals = {}
-                if len(keys) == 1:
-                    # only one field: must be the id of the object
-                    vals[keys[0]] = unicode(child_obj.pk)
-                else:
-                    for field in keys:
-                        if hasattr(child_obj, field):
-                            value = getattr(child_obj, field)
-                            if hasattr(value, 'pk'):
-                                value = value.pk
-                            if value is not None:
-                                vals[field] = unicode(value)
-                if vals:
-                    initial.append(vals)
+            return self._get_instanced_init_for_formset(obj, current_step,
+                                                        c_form)
         return initial
 
 
@@ -1418,6 +1467,7 @@ class AccountWizard(Wizard):
     model = models.Person
     formset_pop_deleted = False
     wizard_done_window = reverse_lazy('show-person')
+    modification = True
 
     def get_formated_datas(self, forms):
         datas = super(AccountWizard, self).get_formated_datas(forms)
@@ -1482,7 +1532,6 @@ class AccountWizard(Wizard):
         profile_form = form_dict['profile-account_management']
         for form in profile_form:
             data = form.cleaned_data
-            print(data)
             profile = None
             if data.get('pk', None):
                 try:
@@ -1522,7 +1571,6 @@ class AccountWizard(Wizard):
             for area in areas:
                 profile.areas.add(area)
 
-
         final_form = form_dict['final-account_management']
         if settings.ADMINS and type(final_form.cleaned_data) == dict and \
                  final_form.cleaned_data.get('send_password', None):