Archaeological files: add departments to preventive forms - Wizards manage several many to many in one form

1 files changed, 23 insertions, 14 deletions

diff --git a/ishtar_common/wizards.py b/ishtar_common/wizards.py
index 7da654b80..cc371b5e2 100644
--- a/ishtar_common/wizards.py
+++ b/ishtar_common/wizards.py
@@ -284,6 +284,8 @@ class Wizard(NamedUrlWizardView):
                     frm = form.forms[0]
                     if hasattr(frm, 'base_model') and frm.base_model:
                         whole_associated_models.append(frm.base_model)
+                    elif hasattr(frm, 'base_models') and frm.base_models:
+                        whole_associated_models += frm.base_models
                     else:
                         whole_associated_models += associated_models.keys()
                     fields = frm.fields.copy()
@@ -291,8 +293,9 @@ class Wizard(NamedUrlWizardView):
                         fields.pop('DELETE')
                     multi = len(fields) > 1
                     if multi:
-                        assert hasattr(frm, 'base_model'), \
-                            u"Must define a base_model for " + \
+                        assert hasattr(frm, 'base_model') or \
+                            hasattr(frm, 'base_models'), \
+                            u"Must define a base_model(s) for " + \
                             unicode(frm.__class__)
                 for frm in form.forms:
                     if not frm.is_valid():
@@ -314,7 +317,11 @@ class Wizard(NamedUrlWizardView):
                         else:
                             m2m.append((key, value))
                     if multi and vals:
-                        m2m.append((frm.base_model, vals))
+                        if hasattr(frm, 'base_models'):
+                            for m in frm.base_models:
+                                m2m.append((frm.base_model, m))
+                        else:
+                            m2m.append((frm.base_model, vals))
             elif type(form.cleaned_data) == dict:
                 for key in form.cleaned_data:
                     if key.startswith('hidden_'):
@@ -336,15 +343,17 @@ class Wizard(NamedUrlWizardView):
                                 value = model.objects.get(pk=value)
                         else:
                             value = None
-                    if hasattr(form, 'base_model') and form.base_model and \
-                       form.base_model == key:
-                        whole_associated_models.append(form.base_model)
+                    if (hasattr(form, 'base_model') and form.base_model and
+                            form.base_model == key) or (
+                            hasattr(form, 'base_models') and
+                            key in form.base_models):
+                        whole_associated_models.append(key)
                         if value:
                             vals = value
                             if type(vals) not in (list, tuple):
                                 vals = [vals]
                             for val in vals:
-                                m2m.append((form.base_model, val))
+                                m2m.append((key, val))
                     else:
                         dct[key] = value
         return self.save_model(dct, m2m, whole_associated_models, form_list,
@@ -445,27 +454,27 @@ class Wizard(NamedUrlWizardView):
             obj.save()
         m2m_items = {}
         for model in whole_associated_models:
-            getattr(obj, model+'s').clear()
+            getattr(obj, model + 's').clear()
         for key, value in m2m:
             if key not in m2m_items:
                 if type(key) == dict:
                     vals = []
-                    for item in getattr(obj, key+'s').all():
+                    for item in getattr(obj, key + 's').all():
                         v = {}
                         for k in value.keys():
                             v[k] = getattr(item, k)
                         vals.append(v)
                     m2m_items[key] = vals
                 else:
-                    m2m_items[key] = getattr(obj, key+'s').all()
+                    m2m_items[key] = getattr(obj, key + 's').all()
             if value not in m2m_items[key]:
                 if type(value) == dict:
-                    model = getattr(obj, key+'s').model
+                    model = getattr(obj, key + 's').model
                     if issubclass(model, models.BaseHistorizedItem):
                         value['history_modifier'] = self.request.user
                     value = model.objects.create(**value)
                     value.save()
-                getattr(obj, key+'s').add(value)
+                getattr(obj, key + 's').add(value)
                 # necessary to manage interaction between models like
                 # material_index management for baseitems
                 obj.save()
@@ -979,7 +988,7 @@ class AccountWizard(Wizard):
                 continue
             for key in form.cleaned_data:
                 if key == 'hidden_password' and form.cleaned_data[key]:
-                    datas[-1][1].append((_("New password"), "*"*8))
+                    datas[-1][1].append((_("New password"), "*" * 8))
         return datas
 
     def done(self, form_list, **kwargs):
@@ -1032,7 +1041,7 @@ class AccountWizard(Wizard):
                 'password': dct['password'],
                 'app_name': app_name,
                 'site': site and site.domain or ""
-                })
+            })
             t = loader.get_template('account_activation_email.txt')
             msg = t.render(context)
             subject = _(u"[%(app_name)s] Account creation/modification") % {