Refactor the draw control (#1666)

* First pass implementation of AbstractDrawControl

* Consolidate reset functionality

* Better `None` checks on properties lookup.

* Fix linter issues, safer property indexing.

* Add docstrings to functions and classes

* Fix remove drawn features bug

---------

Co-authored-by: Sufyan Abbasi <[email protected]>
Co-authored-by: Qiusheng Wu <[email protected]>

Author: sufyanAbbasi

geemap/geemap.py

@@ -36,6 +36,88 @@
 basemaps = Box(xyz_to_leaflet(), frozen_box=True)
 
 
+class MapDrawControl(ipyleaflet.DrawControl, map_widgets.AbstractDrawControl):
+    """ "Implements the AbstractDrawControl for the map."""
+
+    _roi_start = False
+    _roi_end = False
+
+    def __init__(self, host_map, **kwargs):
+        """Initialize the map draw control.
+
+        Args:
+            host_map (geemap.Map): The geemap.Map object that the control will be added to.
+        """
+        super(MapDrawControl, self).__init__(host_map=host_map, **kwargs)
+
+    @property
+    def user_roi(self):
+        """Returns the last drawn geometry.
+
+        Returns:
+            ee.Geometry: Last drawn geometry.
+        """
+        return self.last_geometry
+
+    @property
+    def user_rois(self):
+        """Returns all drawn geometries as an ee.FeatureCollection.
+
+        Returns:
+            ee.FeatureCollection: All drawn geometries.
+        """
+        return self.collection
+
+    # NOTE: Overridden for backwards compatibility, where edited geometries are
+    # added to the layer instead of modified in place. Remove when
+    # https://github.com/jupyter-widgets/ipyleaflet/issues/1119 is fixed to
+    # allow geometry edits to be reflected on the tile layer.
+    def _handle_geometry_edited(self, geo_json):
+        return self._handle_geometry_created(geo_json)
+
+    def _get_synced_geojson_from_draw_control(self):
+        return [data.copy() for data in self.data]
+
+    def _bind_to_draw_control(self):
+        # Handles draw events
+        def handle_draw(_, action, geo_json):
+            try:
+                self._roi_start = True
+                if action == "created":
+                    self._handle_geometry_created(geo_json)
+                elif action == "edited":
+                    self._handle_geometry_edited(geo_json)
+                elif action == "deleted":
+                    self._handle_geometry_deleted(geo_json)
+                self._roi_end = True
+                self._roi_start = False
+            except Exception as e:
+                self.reset(clear_draw_control=False)
+                self._roi_start = False
+                self._roi_end = False
+                print("There was an error creating Earth Engine Feature.")
+                raise Exception(e)
+
+        self.on_draw(handle_draw)
+        # NOTE: Uncomment the following code once
+        # https://github.com/jupyter-widgets/ipyleaflet/issues/1119 is fixed
+        # to allow edited geometries to be reflected instead of added.
+        # def handle_data_update(_):
+        #     self._sync_geometries()
+        # self.observe(handle_data_update, 'data')
+
+    def _remove_geometry_at_index_on_draw_control(self, index):
+        # NOTE: Uncomment the following code once
+        # https://github.com/jupyter-widgets/ipyleaflet/issues/1119 is fixed to
+        # remove drawn geometries with `remove_last_drawn()`.
+        # del self.data[index]
+        # self.send_state(key='data')
+        pass
+
+    def _clear_draw_control(self):
+        return self.clear()
+
+
 class Map(ipyleaflet.Map):
     """The Map class inherits the ipyleaflet Map class. The arguments you can pass to the Map initialization
         can be found at https://ipyleaflet.readthedocs.io/en/latest/map_and_basemaps/map.html.
@@ -46,6 +128,27 @@ class Map(ipyleaflet.Map):
         object: ipyleaflet map object.
     """
 
+    # Map attributes for drawing features
+    @property
+    def draw_features(self):
+        return self.draw_control.features if self.draw_control else []
+
+    @property
+    def draw_last_feature(self):
+        return self.draw_control.last_feature if self.draw_control else None
+
+    @property
+    def draw_layer(self):
+        return self.draw_control.layer if self.draw_control else None
+
+    @property
+    def user_roi(self):
+        return self.draw_control.user_roi if self.draw_control else None
+
+    @property
+    def user_rois(self):
+        return self.draw_control.user_rois if self.draw_control else None
+
     def __init__(self, **kwargs):
         """Initialize a map object. The following additional parameters can be passed in addition to the ipyleaflet.Map parameters:
 
@@ -170,13 +273,6 @@ def __init__(self, **kwargs):
                 if kwargs.get(control, True):
                     self.add_controls(control, position="bottomright")
 
-        # Map attributes for drawing features
-        self.draw_features = []
-        self.draw_last_feature = None
-        self.draw_layer = None
-        self.user_roi = None
-        self.user_rois = None
-
         # Map attributes for layers
         self.geojson_layers = []
         self.ee_layers = []
@@ -2467,7 +2563,9 @@ def _on_close():
                 self.inspector_control.close()
                 self.inspector_control = None
 
-        inspector = map_widgets.Inspector(self, names, visible, decimals, opened, show_close_button)
+        inspector = map_widgets.Inspector(
+            self, names, visible, decimals, opened, show_close_button
+        )
         inspector.on_close = _on_close
         self.inspector_control = ipyleaflet.WidgetControl(
             widget=inspector, position=position
@@ -2500,8 +2598,8 @@ def add_draw_control(self, position="topleft"):
         Args:
             position (str, optional): The position of the draw control. Defaults to "topleft".
         """
-
-        draw_control = ipyleaflet.DrawControl(
+        draw_control = MapDrawControl(
+            host_map=self,
             marker={"shapeOptions": {"color": "#3388ff"}},
             rectangle={"shapeOptions": {"color": "#3388ff"}},
             # circle={"shapeOptions": {"color": "#3388ff"}},
@@ -2510,50 +2608,6 @@ def add_draw_control(self, position="topleft"):
             remove=True,
             position=position,
         )
-
-        # Handles draw events
-        def handle_draw(target, action, geo_json):
-            try:
-                self._roi_start = True
-                geom = geojson_to_ee(geo_json, False)
-                self.user_roi = geom
-                feature = ee.Feature(geom)
-                self.draw_last_feature = feature
-                if not hasattr(self, "_draw_count"):
-                    self._draw_count = 0
-                if action == "deleted" and len(self.draw_features) > 0:
-                    self.draw_features.remove(feature)
-                    self._draw_count -= 1
-                else:
-                    self.draw_features.append(feature)
-                    self._draw_count += 1
-                collection = ee.FeatureCollection(self.draw_features)
-                self.user_rois = collection
-                ee_draw_layer = EELeafletTileLayer(
-                    collection, {"color": "blue"}, "Drawn Features", False, 0.5
-                )
-                draw_layer_index = self.find_layer_index("Drawn Features")
-
-                if draw_layer_index == -1:
-                    self.add(ee_draw_layer)
-                    self.draw_layer = ee_draw_layer
-                else:
-                    self.substitute_layer(self.draw_layer, ee_draw_layer)
-                    self.draw_layer = ee_draw_layer
-                self._roi_end = True
-                self._roi_start = False
-            except Exception as e:
-                self._draw_count = 0
-                self.draw_features = []
-                self.draw_last_feature = None
-                self.draw_layer = None
-                self.user_roi = None
-                self._roi_start = False
-                self._roi_end = False
-                print("There was an error creating Earth Engine Feature.")
-                raise Exception(e)
-
-        draw_control.on_draw(handle_draw)
         self.add(draw_control)
         self.draw_control = draw_control
 
@@ -2585,8 +2639,11 @@ def add_toolbar(self, position="topright", **kwargs):
         """
 
         from .toolbar import Toolbar, main_tools, extra_tools
+
         self._toolbar = Toolbar(self, main_tools, extra_tools)
-        toolbar_control = ipyleaflet.WidgetControl(widget=self._toolbar, position=position)
+        toolbar_control = ipyleaflet.WidgetControl(
+            widget=self._toolbar, position=position
+        )
         self.add(toolbar_control)
 
     def add_plot_gui(self, position="topright", **kwargs):
@@ -4060,46 +4117,37 @@ def add_remote_tile(
         else:
             raise Exception("The source must be a URL.")
 
+    def remove_draw_control(self):
+        """Removes the draw control from the map"""
+        controls = []
+        old_draw_control = None
+        for control in self.controls:
+            if isinstance(control, MapDrawControl):
+                old_draw_control = control
+
+            else:
+                controls.append(control)
+
+        self.controls = tuple(controls)
+        if old_draw_control:
+            old_draw_control.close()
+
     def remove_drawn_features(self):
         """Removes user-drawn geometries from the map"""
-        if self.draw_layer is not None:
-            self.remove_layer(self.draw_layer)
-            self._draw_count = 0
-            self.draw_features = []
-            self.draw_last_feature = None
-            self.draw_layer = None
-            self.user_roi = None
-            self.user_rois = None
-            self._chart_values = []
-            self._chart_points = []
-            self._chart_labels = None
         if self.draw_control is not None:
-            self.draw_control.clear()
+            self.draw_control.reset()
 
     def remove_last_drawn(self):
-        """Removes user-drawn geometries from the map"""
-        if self.draw_layer is not None:
-            collection = ee.FeatureCollection(self.draw_features[:-1])
-            ee_draw_layer = EELeafletTileLayer(
-                collection, {"color": "blue"}, "Drawn Features", True, 0.5
-            )
-            if self._draw_count == 1:
+        """Removes last user-drawn geometry from the map"""
+        if self.draw_control is not None:
+            if self.draw_control.count == 1:
                 self.remove_drawn_features()
-            else:
-                self.substitute_layer(self.draw_layer, ee_draw_layer)
-                self.draw_layer = ee_draw_layer
-                self._draw_count -= 1
-                self.draw_features = self.draw_features[:-1]
-                self.draw_last_feature = self.draw_features[-1]
-                self.draw_layer = ee_draw_layer
-                self.user_roi = ee.Feature(
-                    collection.toList(collection.size()).get(
-                        collection.size().subtract(1)
-                    )
-                ).geometry()
-                self.user_rois = collection
-                self._chart_values = self._chart_values[:-1]
-                self._chart_points = self._chart_points[:-1]
+            elif self.draw_control.count:
+                self.draw_control.remove_geometry(self.draw_control.geometries[-1])
+                if hasattr(self, "_chart_values"):
+                    self._chart_values = self._chart_values[:-1]
+                if hasattr(self, "_chart_points"):
+                    self._chart_points = self._chart_points[:-1]
                 # self._chart_labels = None
 
     def extract_values_to_points(self, filename):