Context record: can change operation (refs #3080)

1 files changed, 18 insertions, 0 deletions

diff --git a/archaeological_context_records/wizards.py b/archaeological_context_records/wizards.py
index 2449b400e..2780acfbc 100644
--- a/archaeological_context_records/wizards.py
+++ b/archaeological_context_records/wizards.py
@@ -52,6 +52,15 @@ class RecordWizard(Wizard):
                 return current_ope
             except(TypeError, ValueError, ObjectDoesNotExist):
                 pass
+        else:
+            ope_form_key = 'operation-' + self.url_name
+            try:
+                idx = int(self.session_get_value(
+                    ope_form_key, 'operation'))
+                current_ope = models.Operation.objects.get(pk=idx)
+                return current_ope
+            except(TypeError, ValueError, ObjectDoesNotExist):
+                pass
         current_cr = self.get_current_object()
         if current_cr:
             return current_cr.operation
@@ -92,6 +101,15 @@ class RecordWizard(Wizard):
                     data['operation'] = current_obj
                 except(TypeError, ValueError, ObjectDoesNotExist):
                     pass
+            elif step.endswith('_modification'):
+                ope_form_key = 'operation-' + self.url_name
+                try:
+                    idx = int(self.session_get_value(ope_form_key,
+                                                     'operation'))
+                    current_obj = models.Operation.objects.get(pk=idx)
+                    data['operation'] = current_obj
+                except(TypeError, ValueError, ObjectDoesNotExist):
+                    pass
             else:
                 current_object = self.get_current_object()
                 data['context_record'] = current_object