From f7eb7230b22e40b69e176b1a74459793e3150388 Mon Sep 17 00:00:00 2001
From: atmorling <atmorling@gmail.com>
Date: Fri, 29 Nov 2024 16:56:18 +0200
Subject: [PATCH] filter NaNs if they arrive in the legend (#335)

---
 ecoscope/mapping/map.py | 22 +++++++++++++++++++---
 tests/test_ecomap.py    | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/ecoscope/mapping/map.py b/ecoscope/mapping/map.py
index e184ea09..020150fb 100644
--- a/ecoscope/mapping/map.py
+++ b/ecoscope/mapping/map.py
@@ -254,15 +254,31 @@ def add_legend(self, labels: list | pd.Series, colors: list | pd.Series, **kwarg
         style: dict
             Additional style params
         """
+        nans = None
         if isinstance(labels, pd.Series):
+            if pd.api.types.is_numeric_dtype(labels):
+                nans = np.argwhere(np.isnan(labels))
             labels = labels.unique().tolist()
         if isinstance(colors, pd.Series):
             colors = colors.unique().tolist()
 
-        labels = [str(label) for label in labels]
-        colors = [color_tuple_to_css(color) if isinstance(color, tuple) else color for color in colors]
+        if len(labels) != len(colors):
+            raise ValueError("Unique label and color values must be of equal number")
 
-        self.add_widget(LegendWidget(labels=labels, colors=colors, **kwargs))
+        if nans is not None:
+            filtered_labels = []
+            filtered_colors = []
+            for index, value in enumerate(labels):
+                if index not in nans:
+                    filtered_labels.append(labels[index])
+                    filtered_colors.append(colors[index])
+            labels = filtered_labels
+            colors = filtered_colors
+
+        widget_labels = [str(label) for label in labels]
+        widget_colors = [color_tuple_to_css(color) if isinstance(color, tuple) else color for color in colors]
+
+        self.add_widget(LegendWidget(labels=widget_labels, colors=widget_colors, **kwargs))
 
     def add_north_arrow(self, **kwargs):
         """
diff --git a/tests/test_ecomap.py b/tests/test_ecomap.py
index 3c740e63..c4955b3c 100644
--- a/tests/test_ecomap.py
+++ b/tests/test_ecomap.py
@@ -2,6 +2,7 @@
 import random
 import ee
 import geopandas as gpd
+import numpy as np
 import pandas as pd
 import pytest
 import ecoscope
@@ -16,6 +17,7 @@
     TitleWidget,
     SaveImageWidget,
     FullscreenWidget,
+    LegendWidget,
 )
 
 
@@ -81,6 +83,37 @@ def test_add_legend_series():
     m = EcoMap(default_widgets=False)
     m.add_legend(labels=pd.Series(["Black", "White"]), colors=pd.Series([(0, 0, 0, 255), (255, 255, 255, 255)]))
     assert len(m.deck_widgets) == 1
+    legend = m.deck_widgets[0]
+    assert isinstance(legend, LegendWidget)
+    assert legend.labels == ["Black", "White"]
+    assert legend.colors == ["rgba(0, 0, 0, 1.0)", "rgba(255, 255, 255, 1.0)"]
+
+
+def test_add_legend_series_with_nan():
+    m = EcoMap(default_widgets=False)
+    m.add_legend(
+        labels=pd.Series([0, 1, np.nan, 5, np.nan]),
+        colors=pd.Series([(0, 0, 0, 255), (255, 255, 255, 255), (0, 0, 0, 0), (100, 100, 100, 255), (0, 0, 0, 0)]),
+    )
+    assert len(m.deck_widgets) == 1
+    legend = m.deck_widgets[0]
+    assert isinstance(legend, LegendWidget)
+    assert legend.labels == ["0.0", "1.0", "5.0"]
+    assert legend.colors == ["rgba(0, 0, 0, 1.0)", "rgba(255, 255, 255, 1.0)", "rgba(100, 100, 100, 1.0)"]
+
+
+def test_add_legend_series_unbalanced_good():
+    m = EcoMap(default_widgets=False)
+    m.add_legend(
+        labels=pd.Series(["Black", "White", "White"]),
+        colors=pd.Series([(0, 0, 0, 255), (255, 255, 255, 255), (255, 255, 255, 255)]),
+    )
+
+
+def test_add_legend_series_unbalanced_bad():
+    m = EcoMap(default_widgets=False)
+    with pytest.raises(ValueError, match="Unique label and color values must be of equal number"):
+        m.add_legend(labels=pd.Series(["Black", "White"]), colors=pd.Series([(0, 0, 0, 255), (0, 0, 0, 255)]))
 
 
 def test_add_north_arrow():