fix raw_data for loc_resolver

wschuell · Nov 23, 2023 · 8a6ee41 · 8a6ee41
1 parent ab4f601
commit 8a6ee41
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/gis_fillers/fillers/loc_resolver.py b/gis_fillers/fillers/loc_resolver.py
@@ -65,8 +65,12 @@ def prepare(self):
             self.location_list
         ), f"Mismatch between results length{len(results)} and expected length {len(self.location_list)}"
         for res, loc in zip(results, self.location_info):
-            loc["geo_lat"] = res[lat_idx]
-            loc["geo_long"] = res[long_idx]
+            if isinstance(res, tuple):
+                loc["geo_lat"] = res[lat_idx]
+                loc["geo_long"] = res[long_idx]
+            else:
+                loc["geo_lat"] = res["lat"]
+                loc["geo_long"] = res["long"]
             for idval, idc in zip(loc["id_list"], self.id_columns):
                 loc[idc] = idval
 

diff --git a/gis_fillers/getters/generic_getters.py b/gis_fillers/getters/generic_getters.py
@@ -17,7 +17,7 @@ def get(self, db, raw_data=False, **kwargs):
         db.cursor.execute(self.query(), self.query_attributes())
         query_result = list(db.cursor.fetchall())
         if raw_data:
-            return query_result
+            return self.parse_results(query_result=query_result)
         else:
             gdf = gpd.GeoDataFrame(
                 self.parse_results(query_result=query_result),
@@ -280,11 +280,14 @@ def parse_results(self, query_result):
             else:
                 geoloc = self.geolocator.geocode(loc)
                 temp_resolved[loc] = geoloc
-            a["geometry"] = shapely.wkt.loads(
-                f"POINT({geoloc.longitude} {geoloc.latitude})"
-            )
-            a["lat"], a["long"] = geoloc.latitude, geoloc.longitude
-            self.fill_cached_address(address=loc, geo_lat=a["lat"], geo_long=a["long"])
+            if geoloc is not None:
+                a["geometry"] = shapely.wkt.loads(
+                    f"POINT({geoloc.longitude} {geoloc.latitude})"
+                )
+                a["lat"], a["long"] = geoloc.latitude, geoloc.longitude
+                self.fill_cached_address(
+                    address=loc, geo_lat=a["lat"], geo_long=a["long"]
+                )
         return ans
 
     def fill_cached_address(self, address, geo_lat, geo_long):