[core] Add colorbar/legend support

geemap/core.py

@@ -575,7 +575,7 @@ def center_object(
 
     def _find_widget_of_type(
         self, widget_type: Type, return_control: bool = False
-    ) -> Optional[Any]:
+    ) -> Optional[ipywidgets.Widget]:
         """Finds a widget in the controls with the passed in type."""
         for widget in self.controls:
             if isinstance(widget, ipyleaflet.WidgetControl):
@@ -807,6 +807,118 @@ def add_layer(
         }
         super().add(tile_layer)
 
+    def _add_legend(
+        self,
+        title: str = "Legend",
+        legend_dict: Dict[str, str] = None,
+        keys: List[Any] = None,
+        colors: List[Any] = None,
+        position: str = "bottomright",
+        builtin_legend: str = None,
+        layer_name: str = None,
+        add_header: bool = True,
+        widget_args: Dict[Any, Any] = None,
+        **kwargs,
+    ):
+        """Adds a customized legend to the map.
+
+        Args:
+            title (str, optional): Title of the legend. Defaults to 'Legend'.
+            legend_dict (dict, optional): A dictionary containing legend items
+                as keys and color as values. If provided, keys and colors will
+                be ignored. Defaults to None.
+            keys (list, optional): A list of legend keys. Defaults to None.
+            colors (list, optional): A list of legend colors. Defaults to None.
+            position (str, optional): Position of the legend. Defaults to
+                'bottomright'.
+            builtin_legend (str, optional): Name of the builtin legend to add
+                to the map. Defaults to None.
+            layer_name (str, optional): The associated layer for the legend.
+                Defaults to None.
+            add_header (bool, optional): Whether the legend can be closed or
+                not. Defaults to True.
+            widget_args (dict, optional): Additional arguments passed to the
+                widget_template() function. Defaults to {}.
+        """
+        legend = map_widgets.Legend(
+            title,
+            legend_dict,
+            keys,
+            colors,
+            position,
+            builtin_legend,
+            add_header,
+            widget_args,
+            **kwargs,
+        )
+        control = ipyleaflet.WidgetControl(widget=legend, position=position)
+        if layer := self.ee_layers.get(layer_name, None):
+            if old_legend := layer.pop("legend", None):
+                self.remove(old_legend)
+            layer["legend"] = control
+
+        super().add(control)
+        return control
+
+    def _add_colorbar(
+        self,
+        vis_params=None,
+        cmap="gray",
+        discrete=False,
+        label=None,
+        orientation="horizontal",
+        position="bottomright",
+        transparent_bg=False,
+        layer_name=None,
+        font_size=9,
+        axis_off=False,
+        max_width=None,
+        **kwargs,
+    ):
+        """Add a matplotlib colorbar to the map.
+
+        Args:
+            vis_params (dict): Visualization parameters as a dictionary. See https://developers.google.com/earth-engine/guides/image_visualization for options.
+            cmap (str, optional): Matplotlib colormap. Defaults to "gray". See https://matplotlib.org/3.3.4/tutorials/colors/colormaps.html#sphx-glr-tutorials-colors-colormaps-py for options.
+            discrete (bool, optional): Whether to create a discrete colorbar. Defaults to False.
+            label (str, optional): Label for the colorbar. Defaults to None.
+            orientation (str, optional): Orientation of the colorbar, such as "vertical" and "horizontal". Defaults to "horizontal".
+            position (str, optional): Position of the colorbar on the map. It can be one of: topleft, topright, bottomleft, and bottomright. Defaults to "bottomright".
+            transparent_bg (bool, optional): Whether to use transparent background. Defaults to False.
+            layer_name (str, optional): The layer name associated with the colorbar. Defaults to None.
+            font_size (int, optional): Font size for the colorbar. Defaults to 9.
+            axis_off (bool, optional): Whether to turn off the axis. Defaults to False.
+            max_width (str, optional): Maximum width of the colorbar in pixels. Defaults to None.
+
+        Raises:
+            TypeError: If the vis_params is not a dictionary.
+            ValueError: If the orientation is not either horizontal or vertical.
+            TypeError: If the provided min value is not scalar type.
+            TypeError: If the provided max value is not scalar type.
+            TypeError: If the provided opacity value is not scalar type.
+            TypeError: If cmap or palette is not provided.
+        """
+        colorbar = map_widgets.Colorbar(
+            vis_params,
+            cmap,
+            discrete,
+            label,
+            orientation,
+            transparent_bg,
+            font_size,
+            axis_off,
+            max_width,
+            **kwargs,
+        )
+        control = ipyleaflet.WidgetControl(widget=colorbar, position=position)
+        if layer := self.ee_layers.get(layer_name, None):
+            if old_colorbar := layer.pop("colorbar", None):
+                self.remove(old_colorbar)
+            layer["colorbar"] = control
+
+        super().add(control)
+        return control
+
     def _open_help_page(
         self, host_map: MapInterface, selected: bool, item: toolbar.Toolbar.Item
     ) -> None:

geemap/geemap.py

@@ -743,38 +743,31 @@ def add_legend(
                 'bottomright'.
             builtin_legend (str, optional): Name of the builtin legend to add
                 to the map. Defaults to None.
+            layer_name (str, optional): The associated layer for the legend.
+                Defaults to None.
             add_header (bool, optional): Whether the legend can be closed or
                 not. Defaults to True.
             widget_args (dict, optional): Additional arguments passed to the
                 widget_template() function. Defaults to {}.
         """
         try:
-            legend = map_widgets.Legend(
+            legend = self._add_legend(
                 title,
                 legend_dict,
                 keys,
                 colors,
                 position,
                 builtin_legend,
+                layer_name,
                 add_header,
                 widget_args,
                 **kwargs,
             )
-
-            legend_control = ipyleaflet.WidgetControl(widget=legend, position=position)
-
-            self._legend_widget = legend
-            self._legend = legend_control
-            self.add(legend_control)
-
+            self._legend = legend
             if not hasattr(self, "legends"):
-                setattr(self, "legends", [legend_control])
+                self.legends = [legend]
             else:
-                self.legends.append(legend_control)
-
-            if layer_name in self.ee_layers:
-                self.ee_layers[layer_name]["legend"] = legend_control
-
+                self.legends.append(legend)
         except Exception as e:
             raise Exception(e)
 
@@ -817,33 +810,25 @@ def add_colorbar(
             TypeError: If cmap or palette is not provided.
         """
 
-        colorbar = map_widgets.Colorbar(
+        colorbar = self._add_colorbar(
             vis_params,
             cmap,
             discrete,
             label,
             orientation,
+            position,
             transparent_bg,
+            layer_name,
             font_size,
             axis_off,
             max_width,
             **kwargs,
         )
-        colormap_ctrl = ipyleaflet.WidgetControl(
-            widget=colorbar, position=position, transparent_bg=transparent_bg
-        )
-
-        self._colorbar = colormap_ctrl
-        if layer_name in self.ee_layers:
-            if "colorbar" in self.ee_layers[layer_name]:
-                self.remove_control(self.ee_layers[layer_name]["colorbar"])
-            self.ee_layers[layer_name]["colorbar"] = colormap_ctrl
+        self._colorbar = colorbar
         if not hasattr(self, "colorbars"):
-            self.colorbars = [colormap_ctrl]
+            self.colorbars = [colorbar]
         else:
-            self.colorbars.append(colormap_ctrl)
-
-        self.add(colormap_ctrl)
+            self.colorbars.append(colorbar)
 
     def remove_colorbar(self):
         """Remove colorbar from the map."""
@@ -852,6 +837,9 @@ def remove_colorbar(self):
 
     def remove_colorbars(self):
         """Remove all colorbars from the map."""
+        for layer in self.ee_layers.values():
+            if widget := layer.pop("colorbar", None):
+                self.remove(widget)
         if hasattr(self, "colorbars"):
             for colorbar in self.colorbars:
                 if colorbar in self.controls:
@@ -865,6 +853,9 @@ def remove_legend(self):
 
     def remove_legends(self):
         """Remove all legends from the map."""
+        for layer in self.ee_layers.values():
+            if widget := layer.pop("legend", None):
+                self.remove(widget)
         if hasattr(self, "legends"):
             for legend in self.legends:
                 if legend in self.controls: