#1742 allowing Choropleth key_on to traverse through array (#1772)

* #1742 allowing Choropleth key_on to traverse through array

* #1742 allowing Choropleth key_on to traverse through array

* adding test case for get_by_key

* fixing pre commit check

* fixes

* avoid repeat code, simplify test

* mypy

---------

Co-authored-by: amrutha1098 <38883175+amrutha1098@users.noreply.github.com>
Co-authored-by: Frank <33519926+Conengmo@users.noreply.github.com>

Author: amrutharajashekar

folium/features.py

@@ -1552,17 +1552,8 @@ def __init__(
 
             key_on = key_on[8:] if key_on.startswith("feature.") else key_on
 
-            def get_by_key(obj, key):
-                return (
-                    obj.get(key, None)
-                    if len(key.split(".")) <= 1
-                    else get_by_key(
-                        obj.get(key.split(".")[0], None), ".".join(key.split(".")[1:])
-                    )
-                )
-
             def color_scale_fun(x):
-                key_of_x = get_by_key(x, key_on)
+                key_of_x = self._get_by_key(x, key_on)
                 if key_of_x is None:
                     raise ValueError(f"key_on `{key_on!r}` not found in GeoJSON.")
 
@@ -1623,6 +1614,20 @@ def highlight_function(x):
         if self.color_scale:
             self.add_child(self.color_scale)
 
+    @classmethod
+    def _get_by_key(cls, obj: Union[dict, list], key: str) -> Union[float, str, None]:
+        key_parts = key.split(".")
+        first_key_part = key_parts[0]
+        if first_key_part.isdigit():
+            value = obj[int(first_key_part)]
+        else:
+            value = obj.get(first_key_part, None)  # type: ignore
+        if len(key_parts) > 1:
+            new_key = ".".join(key_parts[1:])
+            return cls._get_by_key(value, new_key)
+        else:
+            return value
+
     def render(self, **kwargs) -> None:
         """Render the GeoJson/TopoJson and color scale objects."""
         if self.color_scale:

tests/test_features.py

@@ -12,7 +12,7 @@
 from branca.element import Element
 
 import folium
-from folium import ClickForMarker, GeoJson, Map, Popup
+from folium import Choropleth, ClickForMarker, GeoJson, Map, Popup
 
 
 @pytest.fixture
@@ -302,3 +302,27 @@ def test_geometry_collection_get_bounds():
         "type": "GeometryCollection",
     }
     assert folium.GeoJson(geojson_data).get_bounds() == [[0, -3], [4, 2]]
+
+
+def test_choropleth_get_by_key():
+    geojson_data = {
+        "id": "0",
+        "type": "Feature",
+        "properties": {"idx": 0, "value": 78.0},
+        "geometry": {
+            "type": "Polygon",
+            "coordinates": [
+                [
+                    [1, 2],
+                    [3, 4],
+                ]
+            ],
+        },
+    }
+
+    # Test with string path in key_on
+    assert Choropleth._get_by_key(geojson_data, "properties.idx") == 0
+    assert Choropleth._get_by_key(geojson_data, "properties.value") == 78.0
+
+    # Test with combined string path and numerical index in key_on
+    assert Choropleth._get_by_key(geojson_data, "geometry.coordinates.0.0") == [1, 2]