Repository: cogeotiff/rio-cogeo

Issue:
Support rasters which does not fit in memory
As I can see from source code `rio-cogeo` uses `MemoryFile`. Does it mean that this plugin is not suit for cases when source raster is not fit in memory?

Discussion:

Patch:
diff --git a/rio_cogeo/cogeo.py b/rio_cogeo/cogeo.py
index d9e73d2..584e050 100644
--- a/rio_cogeo/cogeo.py
+++ b/rio_cogeo/cogeo.py
@@ -3,6 +3,8 @@
 import os
 import sys
 import warnings
+import tempfile
+from contextlib import contextmanager
 
 import click
 
@@ -27,6 +29,26 @@ from rio_cogeo.utils import (
     _meters_per_pixel,
 )
 
+try:
+    from contextlib import ExitStack
+except ImportError:
+    from contextlib2 import ExitStack
+
+IN_MEMORY_THRESHOLD = int(os.environ.get("IN_MEMORY_THRESHOLD", 10980 * 10980))
+
+
+@contextmanager
+def TemporaryRasterFile(dst_path, suffix=".tif"):
+    """Create temporary file."""
+    fileobj = tempfile.NamedTemporaryFile(
+        dir=os.path.dirname(dst_path), suffix=suffix, delete=False
+    )
+    fileobj.close()
+    try:
+        yield fileobj
+    finally:
+        os.remove(fileobj.name)
+
 
 def cog_translate(
     src_path,
@@ -40,6 +62,7 @@ def cog_translate(
     web_optimized=False,
     latitude_adjustment=True,
     resampling="nearest",
+    in_memory=None,
     config=None,
     quiet=False,
 ):
@@ -71,6 +94,8 @@ def cog_translate(
         Use mercator meters for zoom calculation or ensure max zoom equality.
     resampling : str, optional (default: "nearest")
         Resampling algorithm.
+    in_memory: bool, optional
+        Force processing raster in memory (default: process in memory if small)
     config : dict
         Rasterio Env options.
     quiet: bool, optional (default: False)
@@ -154,62 +179,71 @@ def cog_translate(
                 meta.pop("compress", None)
                 meta.pop("photometric", None)
 
-                with MemoryFile() as memfile:
-                    with memfile.open(**meta) as mem:
-                        wind = list(mem.block_windows(1))
-
-                        if not quiet:
-                            click.echo("Reading input: {}".format(src_path), err=True)
-                        fout = os.devnull if quiet else sys.stderr
-                        with click.progressbar(
-                            wind, length=len(wind), file=fout, show_percent=True
-                        ) as windows:
-                            for ij, w in windows:
-                                matrix = vrt_dst.read(window=w, indexes=indexes)
-                                mem.write(matrix, window=w)
-
-                                if add_mask or mask:
-                                    mask_value = vrt_dst.dataset_mask(window=w)
-                                    mem.write_mask(mask_value, window=w)
-
-                        if not quiet:
-                            click.echo("Adding overviews...", err=True)
-
-                        if overview_level is None:
-                            overview_level = get_maximum_overview_level(
-                                vrt_dst,
-                                min(
-                                    int(dst_kwargs["blockxsize"]),
-                                    int(dst_kwargs["blockysize"]),
-                                ),
-                            )
+                if in_memory is None:
+                    in_memory = vrt_dst.width * vrt_dst.height < IN_MEMORY_THRESHOLD
 
-                        overviews = [2 ** j for j in range(1, overview_level + 1)]
-                        mem.build_overviews(
-                            overviews, ResamplingEnums[overview_resampling]
+                with ExitStack() as ctx:
+                    if in_memory:
+                        tmpfile = ctx.enter_context(MemoryFile())
+                        tmp_dst = ctx.enter_context(tmpfile.open(**meta))
+                    else:
+                        tmpfile = ctx.enter_context(TemporaryRasterFile(dst_path))
+                        tmp_dst = ctx.enter_context(
+                            rasterio.open(tmpfile.name, "w", **meta)
                         )
 
-                        if not quiet:
-                            click.echo("Updating dataset tags...", err=True)
+                    wind = list(tmp_dst.block_windows(1))
+
+                    if not quiet:
+                        click.echo("Reading input: {}".format(src_path), err=True)
+                    fout = os.devnull if quiet else sys.stderr
+                    with click.progressbar(
+                        wind, length=len(wind), file=fout, show_percent=True
+                    ) as windows:
+                        for ij, w in windows:
+                            matrix = vrt_dst.read(window=w, indexes=indexes)
+                            tmp_dst.write(matrix, window=w)
+
+                            if add_mask or mask:
+                                mask_value = vrt_dst.dataset_mask(window=w)
+                                tmp_dst.write_mask(mask_value, window=w)
+
+                    if not quiet:
+                        click.echo("Adding overviews...", err=True)
+
+                    if overview_level is None:
+                        overview_level = get_maximum_overview_level(
+                            vrt_dst,
+                            min(
+                                int(dst_kwargs["blockxsize"]),
+                                int(dst_kwargs["blockysize"]),
+                            ),
+                        )
 
-                        for i, b in enumerate(indexes):
-                            mem.set_band_description(i + 1, src_dst.descriptions[b - 1])
+                    overviews = [2 ** j for j in range(1, overview_level + 1)]
+                    tmp_dst.build_overviews(
+                        overviews, ResamplingEnums[overview_resampling]
+                    )
 
-                        tags = src_dst.tags()
-                        tags.update(
-                            dict(
-                                OVR_RESAMPLING_ALG=ResamplingEnums[
-                                    overview_resampling
-                                ].name.upper()
-                            )
+                    if not quiet:
+                        click.echo("Updating dataset tags...", err=True)
+
+                    for i, b in enumerate(indexes):
+                        tmp_dst.set_band_description(i + 1, src_dst.descriptions[b - 1])
+
+                    tags = src_dst.tags()
+                    tags.update(
+                        dict(
+                            OVR_RESAMPLING_ALG=ResamplingEnums[
+                                overview_resampling
+                            ].name.upper()
                         )
-                        mem.update_tags(**tags)
+                    )
+                    tmp_dst.update_tags(**tags)
 
-                        if not quiet:
-                            click.echo(
-                                "Writing output to: {}".format(dst_path), err=True
-                            )
-                        copy(mem, dst_path, copy_src_overviews=True, **dst_kwargs)
+                    if not quiet:
+                        click.echo("Writing output to: {}".format(dst_path), err=True)
+                    copy(tmp_dst, dst_path, copy_src_overviews=True, **dst_kwargs)
 
 
 def cog_validate(src_path):
diff --git a/rio_cogeo/scripts/cli.py b/rio_cogeo/scripts/cli.py
index 1889570..23f8788 100644
--- a/rio_cogeo/scripts/cli.py
+++ b/rio_cogeo/scripts/cli.py
@@ -12,6 +12,8 @@ from rasterio.enums import Resampling as ResamplingEnums
 from rio_cogeo.cogeo import cog_translate, cog_validate
 from rio_cogeo.profiles import cog_profiles
 
+IN_MEMORY_THRESHOLD = int(os.environ.get("IN_MEMORY_THRESHOLD", 10980 * 10980))
+
 
 class BdxParamType(click.ParamType):
     """Band inddex type."""
@@ -120,6 +122,12 @@ def cogeo():
     ),
     default="nearest",
 )
+@click.option(
+    "--in-memory/--no-in-memory",
+    default=None,
+    help="Force processing raster in memory / not in memory (default: process in memory "
+    "if smaller than {:.0f} million pixels)".format(IN_MEMORY_THRESHOLD // 1e6),
+)
 @click.option("--threads", type=int, default=8)
 @options.creation_options
 @click.option(
@@ -138,6 +146,7 @@ def create(
     web_optimized,
     latitude_adjustment,
     resampling,
+    in_memory,
     threads,
     creation_options,
     quiet,
@@ -172,6 +181,7 @@ def create(
         web_optimized,
         latitude_adjustment,
         resampling,
+        in_memory,
         config,
         quiet,
     )
diff --git a/setup.py b/setup.py
index 08bb985..5b0e4ac 100644
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,9 @@ inst_reqs = [
     "mercantile",
 ]
 
+if sys.version_info < (3, 3):
+    inst_reqs.append("contextlib2")
+
 extra_reqs = {
     "test": ["pytest", "pytest-cov", "rio-tiler"],
     "dev": ["pytest", "pytest-cov", "rio-tiler", "pre-commit"],


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 1394157..46c21df 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -420,6 +420,23 @@ def test_cogeo_validNodataCustom():
         assert result.exit_code == 1
 
 
+def test_cogeo_validTempFile(monkeypatch):
+    """Should work as expected."""
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        result = runner.invoke(
+            cogeo, ["create", raster_path_rgb, "output.tif", "--no-in-memory"]
+        )
+        assert not result.exception
+        assert result.exit_code == 0
+
+    monkeypatch.setenv("IN_MEMORY_THRESHOLD", "100")
+    with runner.isolated_filesystem():
+        result = runner.invoke(cogeo, ["create", raster_path_rgb, "output.tif"])
+        assert not result.exception
+        assert result.exit_code == 0
+
+
 def test_cogeo_validate():
     """Should work as expected."""
     runner = CliRunner()