From 090f2f042570c0d61c3171066031276909cae3a7 Mon Sep 17 00:00:00 2001
From: vincentsarago <vincent.sarago@gmail.com>
Date: Fri, 3 Nov 2023 12:10:00 +0100
Subject: [PATCH] allow sequence and iterator

---
 rio_tiler/mosaic/reader.py | 25 ++++++++++++++++++-------
 rio_tiler/utils.py         | 20 +++++++++++++++++---
 tests/test_mosaic.py       | 14 ++++++++++++++
 3 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/rio_tiler/mosaic/reader.py b/rio_tiler/mosaic/reader.py
index a7d7a41e..10fcb83b 100644
--- a/rio_tiler/mosaic/reader.py
+++ b/rio_tiler/mosaic/reader.py
@@ -2,7 +2,18 @@
 
 import warnings
 from inspect import isclass
-from typing import Any, Callable, List, Optional, Sequence, Tuple, Type, Union, cast, Iterable
+from typing import (
+    Any,
+    Callable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    Union,
+    cast,
+)
 
 import numpy
 from rasterio.crs import CRS
@@ -23,7 +34,7 @@
 
 
 def mosaic_reader(  # noqa: C901
-    mosaic_assets: Iterable,
+    mosaic_assets: Union[Iterator, Sequence],
     reader: Callable[..., ImageData],
     *args: Any,
     pixel_selection: Union[Type[MosaicMethodBase], MosaicMethodBase] = FirstMethod,
@@ -36,7 +47,7 @@ def mosaic_reader(  # noqa: C901
 
     Args:
 
-        mosaic_assets (sequence): List of assets.
+        mosaic_assets (Sequence or Iterator): List of assets.
         reader (callable): Reader function. The function MUST take `(asset, *args, **kwargs)` as arguments, and MUST return an ImageData.
         args (Any): Argument to forward to the reader function.
         pixel_selection (MosaicMethod, optional): Instance of MosaicMethodBase class. Defaults to `rio_tiler.mosaic.methods.defaults.FirstMethod`.
@@ -76,9 +87,10 @@ def mosaic_reader(  # noqa: C901
             "'rio_tiler.mosaic.methods.base.MosaicMethodBase'"
         )
 
-    # if not chunk_size:
-    #    chunk_size = threads if threads > 1 else len(mosaic_assets)
-    chunk_size = threads
+    if not isinstance(mosaic_assets, Iterator) and not chunk_size:
+        chunk_size = threads if threads > 1 else len(mosaic_assets)
+
+    chunk_size = chunk_size or threads
 
     assets_used: List = []
     crs: Optional[CRS]
@@ -86,7 +98,6 @@ def mosaic_reader(  # noqa: C901
     band_names: List[str]
 
     for chunks in _chunks(mosaic_assets, chunk_size):
-        print(threads, len(chunks), chunk_size)
         tasks = create_tasks(reader, chunks, threads, *args, **kwargs)
         for img, asset in filter_tasks(
             tasks,
diff --git a/rio_tiler/utils.py b/rio_tiler/utils.py
index 9bc9e1b5..df1e8ea0 100644
--- a/rio_tiler/utils.py
+++ b/rio_tiler/utils.py
@@ -1,8 +1,20 @@
 """rio_tiler.utils: utility functions."""
 
+import itertools
 import warnings
 from io import BytesIO
-from typing import Any, Dict, Generator, List, Optional, Sequence, Tuple, Union, Iterable
+from typing import (
+    Any,
+    Dict,
+    Generator,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import numpy
 import rasterio
@@ -24,12 +36,14 @@
 from rio_tiler.constants import WEB_MERCATOR_CRS
 from rio_tiler.errors import RioTilerError
 from rio_tiler.types import BBox, ColorMapType, IntervalTuple, RIOResampling
-import itertools
 
 
 def _chunks(my_list: Iterable, chuck_size: int) -> Generator[Sequence, None, None]:
     """Yield successive n-sized chunks from l."""
-    while chunk:= tuple(itertools.islice(my_list, chuck_size)):
+    if not isinstance(my_list, Iterator):
+        my_list = iter(my_list)
+
+    while chunk := tuple(itertools.islice(my_list, chuck_size)):
         yield chunk
 
 
diff --git a/tests/test_mosaic.py b/tests/test_mosaic.py
index 5fb1a2d3..e0e2ca61 100644
--- a/tests/test_mosaic.py
+++ b/tests/test_mosaic.py
@@ -286,6 +286,20 @@ class aClass(object):
     assert m.dtype == "uint8"
 
 
+def test_mosaic_tiler_iter():
+    """Test mosaic tiler with iterator input."""
+    assets_iter = iter(assets)
+
+    (t, m), assets_used = mosaic.mosaic_reader(assets_iter, _read_tile, x, y, z)
+    assert t.shape == (3, 256, 256)
+    assert m.shape == (256, 256)
+    assert m.all()
+    # Should only have value of 1
+    assert numpy.unique(t[0, m == 255]).tolist() == [1]
+    assert t.dtype == "uint16"
+    assert m.dtype == "uint8"
+
+
 def mock_rasterio_open(asset):
     """Mock rasterio Open."""
     assert asset.startswith("http://somewhere-over-the-rainbow.io")