Repository: CartoDB/cartoframes
Pull request #1245 (2019-11-26 11:38:35)

Implement __getitem__ in CDF
Fix https://github.com/CartoDB/cartoframes/issues/1236


Patch:
diff --git a/cartoframes/core/cartodataframe.py b/cartoframes/core/cartodataframe.py
index 6d97421f..5af5ae3c 100644
--- a/cartoframes/core/cartodataframe.py
+++ b/cartoframes/core/cartodataframe.py
@@ -1,19 +1,30 @@
+from pandas import Series
 from geopandas import GeoDataFrame
 
 from ..utils.geom_utils import generate_index, generate_geometry
 
 
 class CartoDataFrame(GeoDataFrame):
-    def __init__(self, *args, **kwargs):
-        """
-        The CartoDataFrame class is an extension of the `geopandas.GeoDataFrame
-        <http://geopandas.org/reference/geopandas.GeoDataFrame.html>`_ class. It provides
-        powerful cartographic visualizations, geometry detection and decoding, and read / write
-        access to the CARTO platform.
-        """
+    """
+    The CartoDataFrame class is an extension of the `geopandas.GeoDataFrame
+    <http://geopandas.org/reference/geopandas.GeoDataFrame.html>`_ class. It provides
+    powerful cartographic visualizations, geometry detection and decoding, and read / write
+    access to the CARTO platform.
+    """
 
+    def __init__(self, *args, **kwargs):
         super(CartoDataFrame, self).__init__(*args, **kwargs)
 
+    def __getitem__(self, key):
+        result = GeoDataFrame.__getitem__(self, key)
+        if isinstance(key, Series):
+            result.__class__ = CartoDataFrame
+        return result
+
+    @property
+    def _constructor(self):
+        return CartoDataFrame
+
     @staticmethod
     def from_carto(*args, **kwargs):
         """
@@ -45,7 +56,6 @@ class CartoDataFrame(GeoDataFrame):
 
                 cdf = CartoDataFrame.from_carto('SELECT * FROM table_name WHERE value > 100')
         """
-
         from ..io.carto import read_carto
         return read_carto(*args, **kwargs)
 
@@ -63,13 +73,9 @@ class CartoDataFrame(GeoDataFrame):
 
                 cdf = CartoDataFrame.from_file('nybb.shp')
         """
-
-        gdf = GeoDataFrame.from_file(filename, **kwargs)
-        return cls(gdf)
-
-    @property
-    def _constructor(self):
-        return CartoDataFrame
+        result = GeoDataFrame.from_file(filename, **kwargs)
+        result.__class__ = cls
+        return result
 
     @classmethod
     def from_features(cls, features, **kwargs):
@@ -85,9 +91,9 @@ class CartoDataFrame(GeoDataFrame):
 
                 cdf = CartoDataFrame.from_features('nybb.shp')
         """
-
-        gdf = GeoDataFrame.from_features(features, **kwargs)
-        return cls(gdf)
+        result = GeoDataFrame.from_features(features, **kwargs)
+        result.__class__ = cls
+        return result
 
     def to_carto(self, *args, **kwargs):
         """
@@ -106,7 +112,6 @@ class CartoDataFrame(GeoDataFrame):
                 cdf = CartoDataFrame.from_file('nybb.shp')
                 cdf.to_carto('table_name', if_exists='replace')
         """
-
         from ..io.carto import to_carto
         return to_carto(self, *args, **kwargs)
 
@@ -173,7 +178,6 @@ class CartoDataFrame(GeoDataFrame):
                 cdf.convert(index_column='my_index')
 
         """
-
         generate_index(self, index_column, drop_index)
         generate_geometry(self, geom_column, lnglat_columns, drop_geom, drop_lnglat)
         return self
@@ -182,6 +186,5 @@ class CartoDataFrame(GeoDataFrame):
         """
         Creates a :py:class:`Map <cartoframes.viz.Map>` visualization
         """
-
         from ..viz import Map, Layer
         return Map(Layer(self, *args, **kwargs))


Test patch:
diff --git a/tests/unit/core/test_cartodataframe.py b/tests/unit/core/test_cartodataframe.py
index c9d5bfca..7ce97cac 100644
--- a/tests/unit/core/test_cartodataframe.py
+++ b/tests/unit/core/test_cartodataframe.py
@@ -1,26 +1,63 @@
 """Unit tests for cartoframes.data."""
 
-import geopandas as gpd
-from cartoframes import CartoDataFrame
 from shapely.geometry import box
+from geopandas import GeoDataFrame
+
+from cartoframes import CartoDataFrame
 
 # DATA FRAME SRC BBOX
 pol_1 = box(1, 1, 2, 2)
 pol_2 = box(3, 3, 4, 4)
-GDF_BOX = gpd.GeoDataFrame({'id': [1, 2], 'geometry': [pol_1, pol_2]}, columns=['id', 'geometry'])
+GDF_BOX = GeoDataFrame({'id': [1, 2], 'geometry': [pol_1, pol_2]}, columns=['id', 'geometry'])
 
 
 class TestCartoDataFrame(object):
 
-    def test_basic_inheritance(self, mocker):
-        """Test basic inheritance"""
+    def test_basic_inheritance(self):
         cdf = CartoDataFrame(GDF_BOX)
         assert isinstance(cdf, CartoDataFrame)
+        assert cdf._constructor == CartoDataFrame
 
-    def test_crs_inhericante(self, mocker):
-        """Test basic inheritance"""
+    def test_crs_inheritance(self):
         cdf = CartoDataFrame(GDF_BOX)
         cdf.crs = 'epsg:4326'
         cdf = cdf.to_crs('epsg:3857')
         assert isinstance(cdf, CartoDataFrame)
         assert cdf.crs == 'epsg:3857'
+
+    def test_filter_inheritance(self):
+        cdf = CartoDataFrame(GDF_BOX)
+        cdf = cdf[cdf.id > 1]
+        assert isinstance(cdf, CartoDataFrame)
+        assert len(cdf) == 1
+
+    def test_from_carto(self, mocker):
+        mock = mocker.patch('cartoframes.io.carto.read_carto')
+        CartoDataFrame.from_carto('table_name')
+        mock.assert_called_once_with('table_name')
+
+    def test_from_file(self, mocker):
+        mock = mocker.patch.object(GeoDataFrame, 'from_file')
+        cdf = CartoDataFrame.from_file('file_name.geojson')
+        mock.assert_called_once_with('file_name.geojson')
+        assert isinstance(cdf, CartoDataFrame)
+
+    def test_from_features(self, mocker):
+        mock = mocker.patch.object(GeoDataFrame, 'from_features')
+        cdf = CartoDataFrame.from_features('__features__')
+        mock.assert_called_once_with('__features__')
+        assert isinstance(cdf, CartoDataFrame)
+
+    def test_to_carto(self, mocker):
+        mock = mocker.patch('cartoframes.io.carto.to_carto')
+        cdf = CartoDataFrame(GDF_BOX)
+        cdf.to_carto('table_name')
+        mock.assert_called_once_with(cdf, 'table_name')
+
+    def test_viz(self, mocker):
+        mocker.patch('cartoframes.viz.Map', return_value='__map__')
+        mock_layer = mocker.patch('cartoframes.viz.Layer')
+        cdf = CartoDataFrame(GDF_BOX)
+        viz = cdf.viz('__style__')
+        mock_layer.assert_called_once_with(cdf, '__style__')
+        assert viz == '__map__'