✨ site profile: ability to set default localisation for treatments

1 files changed, 15 insertions, 3 deletions

diff --git a/archaeological_finds/forms_treatments.py b/archaeological_finds/forms_treatments.py
index be1921568..d5e353a53 100644
--- a/archaeological_finds/forms_treatments.py
+++ b/archaeological_finds/forms_treatments.py
@@ -214,9 +214,22 @@ class BaseTreatmentForm(CustomForm, ManageOldType):
 
     def __init__(self, *args, **kwargs):
         user = kwargs.pop('user')
-        super(BaseTreatmentForm, self).__init__(*args, **kwargs)
+        current_profile = get_current_profile()
+        has_default_location = kwargs.get("data", None) or (
+            "initial" in kwargs and
+            kwargs["initial"].get("location", None))
+        super().__init__(*args, **kwargs)
+
+        if not has_default_location and \
+                current_profile.default_location_for_treatment_id:
+            self.fields["location"].initial = \
+                current_profile.default_location_for_treatment_id
+            if current_profile.default_location_for_treatment.organization:
+                self.fields["organization"].initial = \
+                    current_profile.default_location_for_treatment.organization_id
+
         # set current currency
-        currency = get_current_profile().currency
+        currency = current_profile.currency
         for key in ('estimated_cost', 'quoted_cost', 'realized_cost',
                     'insurance_cost'):
             self.fields[key].label = self.fields[key].label.format(
@@ -261,7 +274,6 @@ class BaseTreatmentForm(CustomForm, ManageOldType):
         except models.TreatmentType.DoesNotExist:
             raise forms.ValidationError(_("Unknow treatment type"))
 
-
         change_current_location = [
             str(tp) for tp in treatment_types
             if tp.change_current_location