Many fix on CSV imports - allow to import only slice of rows

3 files changed, 66 insertions, 28 deletions

diff --git a/archaeological_operations/forms.py b/archaeological_operations/forms.py
index 7cd076dec..e545ff5fa 100644
--- a/archaeological_operations/forms.py
+++ b/archaeological_operations/forms.py
@@ -241,6 +241,8 @@ class OperationFormGeneral(forms.Form):
                                 required=False)
         code_dracar = forms.CharField(label=u"Code DRACAR", required=False,
                            validators=[validators.MaxLengthValidator(10)])
+        eas_number = forms.CharField(label=u"Numéro de l'EA", required=False,
+                           validators=[validators.MaxLengthValidator(20)])
         cira_date = forms.DateField(label=u"Date avis CIRA", required=False,
                                     widget=widgets.JQueryDate)
         cira_rapporteur = forms.IntegerField(label=u"Rapporteur CIRA",
diff --git a/archaeological_operations/import_from_csv.py b/archaeological_operations/import_from_csv.py
index b4e18635a..97d60e67c 100644
--- a/archaeological_operations/import_from_csv.py
+++ b/archaeological_operations/import_from_csv.py
@@ -33,7 +33,7 @@ from django.contrib.auth.models import User
 from django.db import transaction
 
 from ishtar_common.models import Town, Person, PersonType, OrganizationType, \
-    Organization
+    Organization, SourceType
 from archaeological_files.models import PermitType, File, FileType
 from archaeological_operations.models import Operation, OperationType, Period, \
                                     AdministrativeAct, ActType, OperationSource
@@ -161,6 +161,7 @@ def parse_admin_act_typ(value, code, owner):
 
 def parse_fileref(value):
     value = parse_string(value).split('/')[0]
+    value = value.split('.')[0]
     match = re.search('[0-9].[0-9]*', value)
     if not match:
         return None
@@ -461,9 +462,25 @@ for cols in _OPE_COLS:
     else:
         OPE_COLS.append(None)
 
+
+def ope_postimportfix(ope, dct):
+    changed = False
+    if not ope.year:
+        sd = dct.get('start_date')
+        ed = dct.get('end_date')
+        if sd:
+            ope.year = sd.year
+            changed = True
+        elif ed:
+            ope.year = ed.year
+            changed = True
+    if changed:
+        ope.save()
+    return ope
+
 #@transaction.commit_manually
 def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
-                      stdout=None):
+                      stdout=None, lines=None):
     default_person = person or User.objects.order_by('pk').all()[0]
     # key : (class, default, reverse)
     key_classes = {
@@ -481,7 +498,20 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
     error_ope, error_reversed, error_multis = [], [], []
     multi_keys = set([column.col_models[0]
                       for column in col_defs if column and column.multi])
+
+    restrict_lines = []
+    if lines:
+        if '-' not in lines:
+            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))
     for line_idx, vals in enumerate(values):
+        if restrict_lines:
+            if line_idx > max(restrict_lines):
+                break
+            if line_idx not in restrict_lines:
+                continue
         if stdout:
             stdout.write("\r* line %d" % (line_idx))
         if not line_idx:
@@ -502,13 +532,13 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
                 c_args = c_args[attr]
             try:
                 if not extra_cols:
-                    v = typ(val)
+                        v = typ(val)
                 else:
-                    arguments = [vals[col_number] for col_number in extra_cols]
-                    if not [arg for arg in arguments if arg]:
-                        continue
-                    arguments += [default_person]
-                    v = typ(val, *arguments)
+                        arguments = [vals[col_number] for col_number in extra_cols]
+                        if not [arg for arg in arguments if arg]:
+                            continue
+                        arguments += [default_person]
+                        v = typ(val, *arguments)
             except:
                 v = None
             if len(attrs) == 1:
@@ -541,13 +571,13 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
                     args.pop(k)
                     continue
                 args.pop(k)
-                attached_instance_models[k] = default.copy()
+                attached_instance_models[k] = (cls, default.copy())
             elif type(args[k]) == list or k in multi_keys:
                 multis.append((k, args[k]))
                 args.pop(k)
             elif '__' in k:
                 mod, value = k.split('__')
-                attached_models[mod] = args.pop(k)
+                attached_models[(mod, value)] = args.pop(k)
         op = None
         if not update and not args.get('operation_type'):
             #print "Pas d'operation_type"
@@ -567,12 +597,11 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
             args['year'] = args['start_date'].year
         # creation
         """
-        print args
-        print reversed_items
-        print multis
-        print attached_models
-        print attached_instance_models
-        """
+        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'):
@@ -606,22 +635,24 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
         except:
             #transaction.rollback()
             error_reversed.append((line_idx, reversed_items))
-            continue
         try:
             for k, vals in multis:
+                if not vals:
+                    continue
                 for v in vals:
                     getattr(op, k).add(v)
                     op.save()
         except:
             #transaction.rollback()
             error_multis.append((line_idx, multis))
-            continue
-        for attr in attached_models:
-            setattr(op, attr, attached_models[attr])
-            op.save()
+        for attr, attached_attr in attached_models:
+            field = getattr(op, attr)
+            if field:
+                setattr(field, attached_attr, attached_models[(attr, attached_attr)])
+                field.save()
         #transaction.commit()
         for attr in attached_instance_models:
-            default = attached_instance_models[attr]
+            cls, default = attached_instance_models[attr]
             obj = getattr(op, attr)
             if not obj:
                 obj = cls.objects.create(**default)
@@ -631,6 +662,8 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
                 for k in default:
                     setattr(obj, k, default[k])
                 obj.save()
+        ope_postimportfix(op, args)
+        raise
         #transaction.commit()
 
     errors = []
@@ -653,8 +686,8 @@ def import_operations_csv(values, col_defs=OPE_COLS, update=True, person=None,
     #transaction.commit()
     return new_ops, errors
 
-def import_from_csv(filename, update=False, col_defs=OPE_COLS,
-                    person=None, stdout=None):
+def import_from_csv(filename, update=True, col_defs=OPE_COLS,
+                    person=None, stdout=None, lines=None):
     """
     Import from a CSV file.
     Return number of operation treated and errors.
@@ -666,5 +699,6 @@ def import_from_csv(filename, update=False, col_defs=OPE_COLS,
         return 0, [u"Incorrect CSV file."]
 
     new_ops, errors = import_operations_csv(values, col_defs=col_defs,
-                                    update=update, person=person, stdout=stdout)
+                                    update=update, person=person, stdout=stdout,
+                                    lines=lines)
     return new_ops, errors
diff --git a/archaeological_operations/management/commands/import_operations.py b/archaeological_operations/management/commands/import_operations.py
index 3cf4a569d..acc43a591 100755
--- a/archaeological_operations/management/commands/import_operations.py
+++ b/archaeological_operations/management/commands/import_operations.py
@@ -28,7 +28,7 @@ IMPORTERS = {'csv':import_from_csv,
              'vfp':import_from_dbf}
 
 class Command(BaseCommand):
-    args = '<filename> [<update> <csv|dbf>]'
+    args = '<filename> [<update> <csv|dbf> <lines>]'
     help = "Import archaelogical operations"
 
     def handle(self, *args, **options):
@@ -36,7 +36,8 @@ class Command(BaseCommand):
             raise CommandError("No file provided.")
         filename = args[0]
         update = len(args) > 1 and args[1]
-        file_type = len(args) > 1 and args[2]
+        file_type = len(args) > 2 and args[2]
+        lines = len(args) > 3 and args[3]
         if not file_type:
             suffix = filename.split('.')[-1].lower()
             if suffix in IMPORTERS.keys():
@@ -48,7 +49,8 @@ class Command(BaseCommand):
             raise CommandError("This file type is not managed.")
         nb_ops, errors = IMPORTERS[file_type](filename,
                                               update=update,
-                                              stdout=self.stdout)
+                                              stdout=self.stdout,
+                                              lines=lines)
         self.stdout.write('\n* %d operation treated\n' % nb_ops)
         if errors:
             self.stderr.write('\n'.join(errors))