From 15abfb207405d3184356f33425a6f8c07107224e Mon Sep 17 00:00:00 2001
From: Nathaniel Schmitz <schmitznathaniel@gmail.com>
Date: Thu, 28 Sep 2023 07:52:53 -0400
Subject: [PATCH] Fix layer management issues (#1740)

* Listen to layer updates for the layer manager.

* Fix the toggle all checkbox in layer manager.

* Fix map_widgets unit tests.

* Allow the basemap to be turned off

---------

Co-authored-by: Qiusheng Wu <giswqs@gmail.com>
---
 geemap/core.py        | 12 ++++++++----
 geemap/map_widgets.py | 18 ++++++++++++------
 tests/fake_map.py     |  5 +++++
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/geemap/core.py b/geemap/core.py
index 35b6ad88d5..dca0f310d1 100644
--- a/geemap/core.py
+++ b/geemap/core.py
@@ -197,6 +197,9 @@ def __init__(self, **kwargs):
         if kwargs.get("ee_initialize", True):
             common.ee_initialize()
 
+        # Listen for layers being added/removed so we can update the layer manager.
+        self.observe(self._on_layers_change, "layers")
+
     def get_zoom(self) -> int:
         return self.zoom
 
@@ -303,8 +306,6 @@ def add(self, obj: Any, position: str = "topright", **kwargs) -> None:
             self._add_basemap_selector(position, **kwargs)
         else:
             super().add(obj)
-        if self._layer_manager:
-            self._layer_manager.refresh_layers()
 
     def _on_toggle_toolbar_layers(self, is_open: bool) -> None:
         if is_open:
@@ -459,8 +460,6 @@ def add_layer(
             "vis_params": vis_params,
         }
         super().add(tile_layer)
-        if self._layer_manager:
-            self._layer_manager.refresh_layers()
 
     addLayer = add_layer
 
@@ -542,3 +541,8 @@ def _get_preferred_basemap_name(self, basemap_name: str) -> str:
             ]
         except ValueError:
             return basemap_name
+
+    def _on_layers_change(self, change) -> None:
+        del change  # Unused.
+        if self._layer_manager:
+            self._layer_manager.refresh_layers()
diff --git a/geemap/map_widgets.py b/geemap/map_widgets.py
index 303988f359..01995f17fe 100644
--- a/geemap/map_widgets.py
+++ b/geemap/map_widgets.py
@@ -923,10 +923,11 @@ def refresh_layers(self):
         )
         toggle_all_checkbox.observe(self._on_all_layers_visibility_toggled, "value")
 
-        layer_rows = [toggle_all_checkbox]
-        for layer in self._host_map.layers[1:]:  # Skip the basemap.
+        layer_rows = []
+        non_basemap_layers = self._host_map.layers[1:]  # Skip the basemap.
+        for layer in non_basemap_layers:
             layer_rows.append(self._render_layer_row(layer))
-        self._toolbar_footer.children = layer_rows
+        self._toolbar_footer.children = [toggle_all_checkbox] + layer_rows
 
     def _on_close_click(self, _):
         if self.on_close:
@@ -989,9 +990,14 @@ def _on_layer_settings_click(self, button):
             self.on_open_vis(button.tooltip)
 
     def _on_all_layers_visibility_toggled(self, change):
-        for layer in self._host_map.layers[1:]:  # Skip the basemap.
-            if hasattr(layer, "visible"):
-                layer.visible = change["new"]
+        checkboxes = [
+            row.children[0] for row in self._toolbar_footer.children[1:]
+        ]  # Skip the all on/off checkbox.
+        self._host_map.layers[0].visible = change[
+            "new"
+        ]  # The first layer is the basemap.
+        for checkbox in checkboxes:
+            checkbox.value = change["new"]
 
     def _on_layer_opacity_changed(self, change, layer):
         if layer in self._host_map.geojson_layers:
diff --git a/tests/fake_map.py b/tests/fake_map.py
index 2354f90991..648bfd7ce8 100644
--- a/tests/fake_map.py
+++ b/tests/fake_map.py
@@ -9,6 +9,7 @@ def __init__(self):
         self.layers = []
         self.ee_layers = {}
         self.geojson_layers = []
+        self.controls = []
 
         self._recognized_attrs = self.__dict__.keys()
 
@@ -45,6 +46,10 @@ def find_layer_index(self, name):
 
         return -1
 
+    def add(self, obj):
+        del obj  # Unused.
+        pass
+
     def add_layer(self, layer):
         self.layers.append(layer)