diff --git a/folium/utilities.py b/folium/utilities.py index d6f40580b1..e9807ad91d 100644 --- a/folium/utilities.py +++ b/folium/utilities.py @@ -122,6 +122,8 @@ def _validate_locations_basics(locations: TypeMultiLine) -> None: def validate_locations(locations: TypeLine) -> list[list[float]]: """Validate an iterable with lat/lon coordinate pairs.""" locations = if_pandas_df_convert_to_numpy(locations) + if isinstance(locations, Iterator): + locations = list(locations) _validate_locations_basics(locations) return [validate_location(coord_pair) for coord_pair in locations] @@ -131,7 +133,12 @@ def validate_multi_locations( ) -> Union[list[list[float]], list[list[list[float]]]]: """Validate an iterable with possibly nested lists of coordinate pairs.""" locations = if_pandas_df_convert_to_numpy(locations) + if isinstance(locations, Iterator): + locations = list(locations) _validate_locations_basics(locations) + # Keep the nesting check below from consuming a point from an inner iterator. + if isinstance(next(iter(locations)), Iterator): + locations = [list(line) for line in locations] try: float(next(iter(next(iter(next(iter(locations))))))) # type: ignore except (TypeError, StopIteration): diff --git a/tests/test_utilities.py b/tests/test_utilities.py index a1f10be361..26d346ca39 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -80,6 +80,20 @@ def test_validate_locations(locations): assert outcome == [[0.0, 5.0], [1.0, 6.0], [2.0, 7.0]] +@pytest.mark.parametrize("validator", [validate_locations, validate_multi_locations]) +@pytest.mark.parametrize("count", [1, 3]) +@pytest.mark.parametrize("iterator_type", [iter, lambda points: (p for p in points)]) +def test_validate_locations_iterators(validator, count, iterator_type): + points = [[i, i + 5] for i in range(count)] + assert validator(iterator_type(points)) == points + + +@pytest.mark.parametrize("validator", [validate_locations, validate_multi_locations]) +def test_validate_locations_empty_iterator(validator): + with pytest.raises(ValueError, match="Locations is empty"): + validator(iter([])) + + @pytest.mark.parametrize( "locations", [ @@ -91,6 +105,15 @@ def test_validate_multi_locations(locations): assert outcome == [[[0, 5], [1, 6], [2, 7]], [[3, 8], [4, 9]]] +@pytest.mark.parametrize("outer_iterator", [False, True]) +def test_validate_multi_locations_iterators(outer_iterator): + lines = [[[0, 5], [1, 6], [2, 7]], [[3, 8], [4, 9]]] + locations = [iter(line) for line in lines] + if outer_iterator: + locations = iter(locations) + assert validate_multi_locations(locations) == lines + + @pytest.mark.parametrize( "locations", [