Merge pull request #1346 from CartoDB/2302-fix-geocoding-none

Fix geom detection for one-None-item geom column
CartoDB · Dec 12, 2019 · 303f17f · 303f17f
2 parents 0302b3f + 9065d72
commit 303f17f
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 6 deletions.
diff --git a/cartoframes/utils/geom_utils.py b/cartoframes/utils/geom_utils.py
@@ -4,6 +4,8 @@
 import shapely
 import binascii as ba
 
+from geopandas import GeoSeries
+
 
 ENC_SHAPELY = 'shapely'
 ENC_WKB = 'wkb'
@@ -18,9 +20,11 @@
 
 def decode_geometry_column(geom_column):
     if geom_column.size > 0:
-        first_geom = next(item for item in geom_column if item is not None)
-        enc_type = detect_encoding_type(first_geom)
-        return geom_column.apply(lambda g: decode_geometry(g, enc_type))
+        enc_type = None
+        if any(geom_column):
+            first_geom = next(item for item in geom_column if item is not None)
+            enc_type = detect_encoding_type(first_geom)
+        return GeoSeries(geom_column.apply(lambda g: decode_geometry(g, enc_type)))
     else:
         return geom_column
 

diff --git a/tests/unit/utils/test_geom_utils.py b/tests/unit/utils/test_geom_utils.py
@@ -1,13 +1,14 @@
 """Unit tests for cartoframes.data.utils"""
 
+import pandas as pd
 import geopandas as gpd
 
 from shapely.geos import lgeos
-from shapely.geometry import Point
+from shapely.geometry import Point, base
 
 from cartoframes.utils.geom_utils import (ENC_EWKT, ENC_SHAPELY, ENC_WKB,
                                           ENC_WKB_BHEX, ENC_WKB_HEX, ENC_WKT,
-                                          decode_geometry, detect_encoding_type)
+                                          decode_geometry_column, decode_geometry, detect_encoding_type)
 
 
 class TestGeomUtils(object):
@@ -21,12 +22,33 @@ def setup_method(self):
         ]
         self.lng = [0, 10, 20]
         self.lat = [0, 15, 30]
-        self.geometry = gpd.geoseries.GeoSeries([
+        self.geometry = gpd.GeoSeries([
             Point([0, 0]),
             Point([10, 15]),
             Point([20, 30])
         ], name='geometry')
 
+    def test_decode_geometry_column(self):
+        geom = pd.Series(['POINT(0 0)', 'POINT(1 1)'])
+        expected_decoded_geom = gpd.GeoSeries([Point([0, 0]), Point([1, 1])])
+
+        decoded_geom = decode_geometry_column(geom)
+        assert str(decoded_geom) == str(expected_decoded_geom)
+
+    def test_decode_geometry_column_empty(self):
+        geom_empty = gpd.GeoSeries([])
+        expected_decoded_geom = gpd.GeoSeries([])
+
+        decoded_geom = decode_geometry_column(geom_empty)
+        assert str(decoded_geom) == str(expected_decoded_geom)
+
+    def test_decode_geometry_column_none(self):
+        geom_none = gpd.GeoSeries([None, None])
+        expected_decoded_geom = gpd.GeoSeries([base.BaseGeometry(), base.BaseGeometry()])
+
+        decoded_geom = decode_geometry_column(geom_none)
+        assert str(decoded_geom) == str(expected_decoded_geom)
+
     def test_detect_encoding_type_shapely(self):
         enc_type = detect_encoding_type(Point(1234, 5789))
         assert enc_type == ENC_SHAPELY