Add wms legend method to Map class (opengeos#608)

* added add_wms_legend method to map class * added test_add_wms_legend unit test * updated stringdoc with url format if using wms legend
sthagen · Nov 18, 2023 · c4bc239 · c4bc239
1 parent 0ffe4a4
commit c4bc239
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/leafmap/foliumap.py b/leafmap/foliumap.py
@@ -417,6 +417,41 @@ def add_wms_layer(
             ).add_to(self)
         except Exception as e:
             raise Exception(e)
+
+    def add_wms_legend(
+        self,
+        url,
+    ):
+        """Add a WMS legend based on an image URL
+
+        Args:
+            url (str): URL of the WMS legend image. Should have this format if using wms legend: {geoserver}/wms?REQUEST=GetLegendGraphic&FORMAT=image/png&LAYER={layer}
+        """
+        from branca.element import Figure, MacroElement, Element
+
+        # Check if the map is a Folium Map instance
+        if not isinstance(self, Map):
+            raise ValueError("The self argument must be an instance of folium.Map.")
+
+        # HTML template for the legend
+        legend_html = f"""
+            {{% macro html(this, kwargs) %}}
+            
+            <div id="maplegend" style="position: fixed; 
+                        bottom: 50px;
+                        right: 50px;
+                        z-index:9999; 
+                        ">
+                <img src="{ url }" alt="legend" style="width: 100%; height: 100%;">
+            </div>
+            {{% endmacro %}}
+        """
+
+        # Create an Element with the HTML and add it to the map
+        macro = MacroElement()
+        macro._template = Template(legend_html)
+
+        self.get_root().add_child(macro)
 
     def add_tile_layer(
         self,

diff --git a/tests/test_foliumap.py b/tests/test_foliumap.py
@@ -153,6 +153,16 @@ def test_add_legend(self):
         out_str = m.to_html()
         assert "NLCD" in out_str
 
+    def test_add_wms_legend(self):
+        """Check add of wms legend based on url"""
+        m = leafmap.Map()
+        url="https://www.mrlc.gov/geoserver/mrlc_display/NLCD_2016_Land_Cover_L48/wms?REQUEST=GetLegendGraphic&FORMAT=image/png&LAYER=NLCD_2016_Land_Cover_L48"
+        m.add_wms_legend(url=url)
+        out_str = m.to_html()
+
+        assert "GetLegendGraphic" in out_str
+
+
     # def test_add_marker_cluster(self):
     #     """Check marker cluster"""
     #     with self.assertRaises(NotImplementedError):