Importing and serializing company data

jarbas/core/management/commands/companies.py

@@ -1,13 +1,23 @@
-import csv
 import lzma
-
-from django.core.exceptions import ValidationError
-from django.core.validators import validate_email
+import rows
 
 from jarbas.core.management.commands import LoadCommand
 from jarbas.core.models import Activity, Company
 
 
+class CompaniesDate(rows.fields.DateField):
+    INPUT_FORMAT = '%d/%m/%Y'
+
+companies_csv_field_types = {
+            'email': rows.fields.EmailField,
+            'opening': CompaniesDate,
+            'situation_date': CompaniesDate,
+            'special_situation_date': CompaniesDate,
+            'latitude': rows.fields.FloatField,
+            'longitude': rows.fields.FloatField
+        }
+
+
 class Command(LoadCommand):
     help = 'Load Serenata de Amor companies dataset into the database'
 
@@ -31,11 +41,12 @@ def save_companies(self):
         skip = ('main_activity', 'secondary_activity')
         keys = tuple(f.name for f in Company._meta.fields if f not in skip)
         with lzma.open(self.path, mode='rt', encoding='utf-8') as file_handler:
-            for row in csv.DictReader(file_handler):
+            for row in rows.import_from_csv(file_handler, force_types=companies_csv_field_types):
+                row = row._asdict()
                 main, secondary = self.save_activities(row)
 
                 filtered = {k: v for k, v in row.items() if k in keys}
-                obj = Company.objects.create(**self.serialize(filtered))
+                obj = Company.objects.create(**filtered)
                 for activity in main:
                     obj.main_activity.add(activity)
                 for activity in secondary:
@@ -62,25 +73,3 @@ def save_activities(self, row):
                 secondaries.append(obj)
 
         return [main], secondaries
-
-    def serialize(self, row):
-        row['email'] = self.to_email(row['email'])
-
-        dates = ('opening', 'situation_date', 'special_situation_date')
-        for key in dates:
-            row[key] = self.to_date(row[key])
-
-        decimals = ('latitude', 'longitude')
-        for key in decimals:
-            row[key] = self.to_number(row[key])
-
-        return row
-
-    @staticmethod
-    def to_email(email):
-        try:
-            validate_email(email)
-            return email
-
-        except ValidationError:
-            return None