Add country-iso2 argument to import admin2

Author: Rub21

api/management/commands/import-admin2-data.py

@@ -14,7 +14,7 @@
 from api.models import Admin2Geoms
 
 class Command(BaseCommand):
-  help = "import a shapefile of administrative boundary level 2 data to the GO database. To run, python manage.py import-admin2-data input.shp"
+  help = "import a shapefile of administrative boundary level 2 data to the GO database. To run, python manage.py import-admin2-data input.shp --country-iso2=af"
 
   missing_args_message = "Filename is missing. A shapefile with valid admin polygons is required."
 
@@ -44,11 +44,16 @@ def add_arguments(self, parser):
       action='store_true',
       help='Import all  admin2 boundaries in the shapefile, if possible.'
     )
+    parser.add_argument(
+      '--country-iso2',
+      type=str,
+      required=True,
+      help='Country iso2 code'
+    )
 
   @transaction.atomic
   def handle(self, *args, **options):
     filename = options['filename'][0]
-  
     # a dict to hold all the admin2 that needs to be manually imported
     import_missing = {}
     if options['import_missing']:
@@ -85,13 +90,10 @@ def handle(self, *args, **options):
 
       centroid = geom.centroid.wkt
       bbox = geom.envelope.wkt
-      # if there is no code, but import all is flagged then import this admin2
-      if code is None or code == '':
-        if options['import_all']:
+      # import all shapes for admin2
+      if options['import_all']:
           self.add_admin2(options, 'all', feature, geom, centroid, bbox)
-
-      # for features that has a code and not NA
-      if code and code != 'N.A':
+      else:
         admin2_objects = Admin2.objects.filter(code=code)
         if len(admin2_objects) == 0:
           if options['import_missing']:
@@ -132,48 +134,49 @@ def handle(self, *args, **options):
             d.bbox = bbox
           d.save()
 
-      if code == 'N.A':
-        admin2 = Admin2.objects.filter(name__icontains=name)
-        if len(admin2):
-          d = admin2[0]
-          if options['update_geom']:
-            self.update_geom(d, geom)
-          if options['update_centroid']:
-            d.centroid = centroid
-          if options['update_bbox']:
-            d.bbox = bbox
-          d.save()
-        else:
-          if options['import_missing']:
-            self.add_admin2(options, import_missing, feature, geom, centroid, bbox)
-          else:
-            missing_file.writerow({'code': code, 'name': name})
-
-      if code == 'f':
-        admin2 = Admin2.objects.filter(name__icontains=name)
-        if not len(admin2):
-          if options['import_missing']:
-            self.add_admin2(options, import_missing, feature, geom, centroid, bbox)
-          else:
-            missing_file.writerow({'code': code, 'name': name})
-
+      # if code == 'N.A':
+      #   admin2 = Admin2.objects.filter(name__icontains=name)
+      #   if len(admin2):
+      #     d = admin2[0]
+      #     if options['update_geom']:
+      #       self.update_geom(d, geom)
+      #     if options['update_centroid']:
+      #       d.centroid = centroid
+      #     if options['update_bbox']:
+      #       d.bbox = bbox
+      #     d.save()
+      #   else:
+      #     if options['import_missing']:
+      #       self.add_admin2(options, import_missing, feature, geom, centroid, bbox)
+      #     else:
+      #       missing_file.writerow({'code': code, 'name': name})
+
+      # if code == 'f':
+      #   admin2 = Admin2.objects.filter(name__icontains=name)
+      #   if not len(admin2):
+      #     if options['import_missing']:
+      #       self.add_admin2(options, import_missing, feature, geom, centroid, bbox)
+      #     else:
+      #       missing_file.writerow({'code': code, 'name': name})
     print('done!')
 
 
   def add_admin2(self, options, import_missing, feature, geom, centroid, bbox):
     code = feature.get('code') or 'N.A'
     name = feature.get('name')
-    admin2 =Admin2()
+    admin2 = Admin2()
     admin2.code = code
     admin2.name = name
     admin2.centroid = centroid
     admin2.bbox = bbox
     try:
       # find district_id based on centroid of admin2 and country.
-      admin1_id = self.find_district_id(centroid)
+      admin1_id = self.find_district_id(centroid, options['country_iso2'])
       if admin1_id is None:
+        print('country does not exist', "af")
         pass
-      admin2.admin1_id = admin1_id
+      else:
+        admin2.admin1_id = admin1_id
     except ObjectDoesNotExist:
       print('country does not exist', "af")
       pass
@@ -195,13 +198,14 @@ def update_geom(self, admin2, geom):
         Admin2Geom.geom = geom
         Admin2Geom.save()
 
-  def find_district_id(self, centroid):
+  def find_district_id(self, centroid, country_iso2):
       """Find district_id for admin2, according to the point with in the district polygon.
       Args:
           centroid (str): Admin2 centroid
+          country_iso2 (str): Country iso2
       """
       admin1_id = None
-      country_id = Country.objects.get(iso="af")
+      country_id = Country.objects.get(iso=country_iso2)
       districts = District.objects.filter(country_id=country_id)
       districts_ids = [d.id for d in districts]
       districts_geoms = DistrictGeoms.objects.filter(district_id__in=districts_ids)