Operation imports: simplify and fix duplicate

2 files changed, 68 insertions, 42 deletions

diff --git a/archaeological_operations/import_from_csv.py b/archaeological_operations/import_from_csv.py
index b2f3e1578..9d3a58ffd 100644
--- a/archaeological_operations/import_from_csv.py
+++ b/archaeological_operations/import_from_csv.py
@@ -478,19 +478,60 @@ def ope_postimportfix(ope, dct):
         ope.save()
     return ope
 
+class RelatedClass:
+    def __init__(self, key, cls, default_data={}, reverse_key=False, unique_keys=[]):
+        self.key, self.cls, self.default_data = key, cls, default_data
+        self.reverse_key, self.unique_keys = reverse_key, unique_keys
+
+    def create_object(self, data):
+        if self.unique_keys:
+            unique_data = {}
+            for k in self.unique_keys:
+                unique_data[k] = data.pop(k)
+            unique_data['defaults'] = data
+            obj, created = self.cls.objects.get_or_create(**unique_data)
+        else:
+            obj = self.cls.objects.create(**data)
+        return obj
+
+    def create(self, item, data, attr=None):
+        if self.reverse_key:
+            data[self.reverse_key] = item
+            obj = self.create_object(data)
+        else:
+            obj = getattr(item, attr)
+            if not obj:
+                obj = self.create_object(data)
+                setattr(item, attr, obj)
+                item.save()
+            else:
+                for k in data:
+                    setattr(obj, k, data[k])
+                obj.save()
+        return obj
+
 #@transaction.commit_manually
 def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
                       stdout=None, lines=None):
     default_person = person or User.objects.order_by('pk').all()[0]
-    # key : (class, default, reverse)
-    key_classes = {
-    'administrative_act':(AdministrativeAct, {'history_modifier':default_person,
-                           'act_type':ActType.objects.get(
-                                      txt_idx='excavation_order')}, 'operation'),
-    'associated_file':(File, {'history_modifier':default_person,
-                              'file_type':FileType.objects.get(txt_idx='undefined')}, None),
-    'source':(OperationSource, {}, 'operation')
-    }
+    RELATED_CLASSES = [
+        RelatedClass('administrative_act', AdministrativeAct,
+                     default_data={'history_modifier':default_person,
+                                     'act_type':ActType.objects.get(
+                                            txt_idx='excavation_order')
+                                     },
+                     reverse_key='operation',
+                     unique_keys=['ref_sra']),
+        RelatedClass('associated_file', File,
+                     default_data={'history_modifier':default_person,
+                                   'file_type':FileType.objects.get(
+                                                txt_idx='undefined')}),
+        RelatedClass('source', OperationSource, reverse_key='operation',
+                     unique_keys=['index']),
+    ]
+    RELATED_CLASSES_KEYS = dict([(rel_cls.key, rel_cls)
+                                   for rel_cls in RELATED_CLASSES])
+    related_classes = RELATED_CLASSES_KEYS
     _prepare_ope_types()
 
     ope_default = {'history_modifier':default_person}
@@ -505,7 +546,7 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
             restrict_lines = [int(line) for line in lines.split(',')]
         else:
             start_line, end_line = lines.split('-')
-            restrict_lines = list(xrange(start_line, end_line+1))
+            restrict_lines = list(xrange(int(start_line), int(end_line)+1))
     for line_idx, vals in enumerate(values):
         if restrict_lines:
             if line_idx > max(restrict_lines):
@@ -560,18 +601,17 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
                 args.pop(k)
                 args.pop(key+'__month')
                 args.pop(key+'__day')
-        reversed_items, multis = [], []
-        attached_models, attached_instance_models = {}, {}
+        related_items = []
+        multis = []
+        attached_models = {}
         for k in args.keys():
-            if k in key_classes:
-                cls, default, reverse = key_classes[k]
+            if k in related_classes:
+                rel_cls = related_classes[k]
+                cls, default = rel_cls.cls, rel_cls.default_data
+                reverse_key = rel_cls.reverse_key
                 default.update(args[k])
-                if reverse:
-                    reversed_items.append((cls, default.copy(), reverse))
-                    args.pop(k)
-                    continue
                 args.pop(k)
-                attached_instance_models[k] = (cls, default.copy())
+                related_items.append((rel_cls, default.copy(), k))
             elif type(args[k]) == list or k in multi_keys:
                 multis.append((k, args[k]))
                 args.pop(k)
@@ -598,10 +638,9 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
         # creation
         """
         print "args", args
-        print "reversed_items", reversed_items
         print "multis", multis
         print "attached_models", attached_models
-        print "attached_instance_models", attached_instance_models"""
+        """
         if not op:
             args.update(ope_default)
             #if not args.get('operation_code'):
@@ -628,14 +667,6 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
                 continue
             #transaction.commit()
         try:
-            for cls, default, reverse in reversed_items:
-                default[reverse] = op
-                it = cls(**default).save()
-            #transaction.commit()
-        except:
-            #transaction.rollback()
-            error_reversed.append((line_idx, reversed_items))
-        try:
             for k, vals in multis:
                 if not vals:
                     continue
@@ -648,20 +679,15 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
         for attr, attached_attr in attached_models:
             field = getattr(op, attr)
             if field:
-                setattr(field, attached_attr, attached_models[(attr, attached_attr)])
+                setattr(field, attached_attr, attached_models[(attr,
+                                                               attached_attr)])
                 field.save()
         #transaction.commit()
-        for attr in attached_instance_models:
-            cls, default = attached_instance_models[attr]
-            obj = getattr(op, attr)
-            if not obj:
-                obj = cls.objects.create(**default)
-                obj.save()
-                setattr(op, attr, obj)
-            else:
-                for k in default:
-                    setattr(obj, k, default[k])
-                obj.save()
+        for rel_cls, data, attr in related_items:
+            try:
+                rel_cls.create(op, data, attr)
+            except:
+                error_reversed.append((line_idx, data))
         ope_postimportfix(op, args)
         #transaction.commit()
 
diff --git a/archaeological_operations/management/commands/import_operations.py b/archaeological_operations/management/commands/import_operations.py
index acc43a591..0c842184b 100755
--- a/archaeological_operations/management/commands/import_operations.py
+++ b/archaeological_operations/management/commands/import_operations.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# Copyright (C) 2012  Étienne Loks <etienne.loks_AT_peacefrogsDOTnet>
+# Copyright (C) 2012-2013  Étienne Loks <etienne.loks_AT_peacefrogsDOTnet>
 
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License as