rio-pmtiles python package [#338] (#542)

Add rio-pmtiles command line tool. [#338] This is derived from the original mapbox/rio-mbtiles implementation, with changes: * output PMTiles only instead of MBTiles. * Python 3.7+ only. * remove --implementation, --image-dump, --append/--overwrite, --covers features. * bump dependency versions. * better progress reporting; add pyroaring. * update README and license texts. * rio-pmtiles v0.0.6 on PyPI
2026-02-04 02:41:09 +00:00 · 2025-03-24 20:50:53 +08:00
parent 1d897f4f7e
commit 63182e525d
36 changed files with 1296 additions and 2 deletions
--- a/python/rio-pmtiles/tests/test_cli.py
+++ b/python/rio-pmtiles/tests/test_cli.py
@@ -0,0 +1,290 @@
+import os
+import sqlite3
+import sys
+import warnings
+
+import click
+from click.testing import CliRunner
+import pytest
+import rasterio
+from rasterio.rio.main import main_group
+
+from pmtiles.reader import Reader, MmapSource, all_tiles
+import rio_pmtiles.scripts.cli
+
+from conftest import mock
+
+class Output:
+    def __init__(self, fname):
+        self.file = open(fname, 'rb')
+
+    def __enter__(self):
+        return Reader(MmapSource(self.file))
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.file.close()
+
+def test_cli_help():
+    runner = CliRunner()
+    result = runner.invoke(main_group, ["pmtiles", "--help"])
+    assert result.exit_code == 0
+    assert "Export a dataset to PMTiles." in result.output
+
+
+@mock.patch("rio_pmtiles.scripts.cli.rasterio")
+def test_dst_nodata_validation(rio):
+    """--dst-nodata requires source nodata in some form"""
+    rio.open.return_value.__enter__.return_value.profile.get.return_value = None
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group, ["pmtiles", "--dst-nodata", "0", "in.tif", "out.pmtiles"]
+    )
+    assert result.exit_code == 2
+
+
+@pytest.mark.parametrize("filename", ["RGB.byte.tif", "RGBA.byte.tif"])
+def test_export_metadata(tmpdir, data, filename):
+    inputfile = str(data.join(filename))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(main_group, ["pmtiles", inputfile, outputfile])
+    assert result.exit_code == 0
+
+    with Output(outputfile) as p:
+        assert p.metadata()['name'] == filename
+
+def test_export_metadata_output_opt(tmpdir, data):
+    inputfile = str(data.join("RGB.byte.tif"))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(main_group, ["pmtiles", inputfile, "-o", outputfile])
+    assert result.exit_code == 0
+
+    with Output(outputfile) as p:
+        assert p.metadata()['name'] == "RGB.byte.tif"
+
+def test_export_tiles(tmpdir, data):
+    inputfile = str(data.join("RGB.byte.tif"))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(main_group, ["pmtiles", inputfile, outputfile])
+    assert result.exit_code == 0
+
+    with open(outputfile, 'rb') as f:
+        src = MmapSource(f)
+        assert len(list(all_tiles(src))) == 6
+
+def test_export_zoom(tmpdir, data):
+    inputfile = str(data.join("RGB.byte.tif"))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group, ["pmtiles", inputfile, outputfile, "--zoom-levels", "6..7"]
+    )
+    assert result.exit_code == 0
+
+    with open(outputfile, 'rb') as f:
+        src = MmapSource(f)
+        assert len(list(all_tiles(src))) == 6
+
+def test_export_jobs(tmpdir, data):
+    inputfile = str(data.join("RGB.byte.tif"))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(main_group, ["pmtiles", inputfile, outputfile, "-j", "4"])
+    assert result.exit_code == 0
+
+    with open(outputfile, 'rb') as f:
+        src = MmapSource(f)
+        assert len(list(all_tiles(src))) == 6
+
+
+def test_export_src_nodata(tmpdir, data):
+    inputfile = str(data.join("RGB.byte.tif"))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group,
+        ["pmtiles", inputfile, outputfile, "--src-nodata", "0", "--dst-nodata", "0"],
+    )
+    assert result.exit_code == 0
+
+    with open(outputfile, 'rb') as f:
+        src = MmapSource(f)
+        assert len(list(all_tiles(src))) == 6
+
+
+def test_export_bilinear(tmpdir, data):
+    inputfile = str(data.join("RGB.byte.tif"))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group, ["pmtiles", inputfile, outputfile, "--resampling", "bilinear"]
+    )
+    assert result.exit_code == 0
+
+    with open(outputfile, 'rb') as f:
+        src = MmapSource(f)
+        assert len(list(all_tiles(src))) == 6
+
+
+def test_skip_empty(tmpdir, empty_data):
+    """This file has the same shape as RGB.byte.tif, but no data."""
+    inputfile = empty_data
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(main_group, ["pmtiles", inputfile, outputfile])
+    assert result.exit_code == 0
+
+    with open(outputfile, 'rb') as f:
+        src = MmapSource(f)
+        assert len(list(all_tiles(src))) == 0
+
+
+def test_include_empty(tmpdir, empty_data):
+    """This file has the same shape as RGB.byte.tif, but no data."""
+    inputfile = empty_data
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group, ["pmtiles", "--include-empty-tiles", inputfile, outputfile]
+    )
+    assert result.exit_code == 0
+
+    with open(outputfile, 'rb') as f:
+        src = MmapSource(f)
+        assert len(list(all_tiles(src))) == 6
+
+
+def test_invalid_format_rgba(tmpdir, empty_data):
+    """--format JPEG --rgba is not allowed"""
+    inputfile = empty_data
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group, ["pmtiles", "--format", "JPEG", "--rgba", inputfile, outputfile]
+    )
+    assert result.exit_code == 2
+
+
+@pytest.mark.parametrize("filename", ["RGBA.byte.tif", "RGB.byte.tif"])
+def test_rgba_png(tmpdir, data, filename):
+    inputfile = str(data.join(filename))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group, ["pmtiles", "--rgba", "--format", "PNG", inputfile, outputfile]
+    )
+    with Output(outputfile) as p:
+        assert p.metadata()['name'] == filename
+
+
+@pytest.mark.parametrize(
+    "minzoom,maxzoom,exp_num_tiles,source",
+    [
+        (4, 10, 70, "RGB.byte.tif"),
+        (6, 7, 6, "RGB.byte.tif"),
+        (4, 10, 12, "rgb-193f513.vrt"),
+        (4, 10, 70, "rgb-fa48952.vrt"),
+    ],
+)
+def test_export_count(tmpdir, data, minzoom, maxzoom, exp_num_tiles, source):
+    inputfile = str(data.join(source))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group,
+        [
+            "pmtiles",
+            "--zoom-levels",
+            "{}..{}".format(minzoom, maxzoom),
+            inputfile,
+            outputfile,
+        ],
+    )
+    assert result.exit_code == 0
+    with open(outputfile, 'rb') as f:
+        src = MmapSource(f)
+        assert len(list(all_tiles(src))) == exp_num_tiles
+
+
+@pytest.mark.parametrize("filename", ["RGBA.byte.tif"])
+def test_progress_bar(tmpdir, data, filename):
+    inputfile = str(data.join(filename))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group,
+        [
+            "pmtiles",
+            "-#",
+            "--zoom-levels",
+            "4..11",
+            "--rgba",
+            "--format",
+            "PNG",
+            inputfile,
+            outputfile,
+        ],
+    )
+    assert result.exit_code == 0
+    assert "100%" in result.output
+
+
+@pytest.mark.parametrize("inputfiles", [[], ["a.tif", "b.tif"]])
+def test_input_required(inputfiles):
+    """We require exactly one input file"""
+    runner = CliRunner()
+    result = runner.invoke(main_group, ["pmtiles"] + inputfiles + ["foo.pmtiles"])
+    assert result.exit_code == 2
+
+
+@pytest.mark.parametrize("filename", ["RGBA.byte.tif"])
+def test_cutline_progress_bar(tmpdir, data, rgba_cutline_path, filename):
+    """rio-pmtiles accepts and uses a cutline"""
+    inputfile = str(data.join(filename))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group,
+        [
+            "pmtiles",
+            "-#",
+            "--zoom-levels",
+            "4..11",
+            "--rgba",
+            "--format",
+            "PNG",
+            "--cutline",
+            rgba_cutline_path,
+            inputfile,
+            outputfile,
+        ],
+    )
+    assert result.exit_code == 0
+    assert "100%" in result.output
+
+
+@pytest.mark.parametrize("filename", ["RGBA.byte.tif"])
+def test_invalid_cutline(tmpdir, data, rgba_points_path, filename):
+    """Points cannot serve as a cutline"""
+    inputfile = str(data.join(filename))
+    outputfile = str(tmpdir.join("export.pmtiles"))
+    runner = CliRunner()
+    result = runner.invoke(
+        main_group,
+        [
+            "pmtiles",
+            "-#",
+            "--zoom-levels",
+            "4..11",
+            "--rgba",
+            "--format",
+            "PNG",
+            "--cutline",
+            rgba_points_path,
+            inputfile,
+            outputfile,
+        ],
+    )
+    assert result.exit_code == 1