From ea1657d8febf219608591eadf5ada266d8da7591 Mon Sep 17 00:00:00 2001 From: Gabor Simko Date: Fri, 9 Feb 2024 20:40:06 -0700 Subject: [PATCH] fix reader.py reading all the tiles (#357) * fix reader.py reading all the tiles * add unit test --------- Co-authored-by: Gabor Simko --- python/pmtiles/reader.py | 7 ++++--- python/test/test_reader_writer.py | 24 +++++++++++++++++++++++- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/python/pmtiles/reader.py b/python/pmtiles/reader.py index 5354107..d0dee2d 100644 --- a/python/pmtiles/reader.py +++ b/python/pmtiles/reader.py @@ -63,9 +63,10 @@ def traverse(get_bytes, header, dir_offset, dir_length): entries = deserialize_directory(get_bytes(dir_offset, dir_length)) for entry in entries: if entry.run_length > 0: - yield tileid_to_zxy(entry.tile_id), get_bytes( - header["tile_data_offset"] + entry.offset, entry.length - ) + for i in range(entry.run_length): + yield tileid_to_zxy(entry.tile_id + i), get_bytes( + header["tile_data_offset"] + entry.offset, entry.length + ) else: for t in traverse( get_bytes, diff --git a/python/test/test_reader_writer.py b/python/test/test_reader_writer.py index ec7647b..7837e58 100644 --- a/python/test/test_reader_writer.py +++ b/python/test/test_reader_writer.py @@ -1,7 +1,7 @@ import unittest from io import BytesIO from pmtiles.writer import Writer -from pmtiles.reader import Reader, MemorySource +from pmtiles.reader import all_tiles, Reader, MemorySource from pmtiles.tile import Compression, TileType, tileid_to_zxy, zxy_to_tileid @@ -46,3 +46,25 @@ class TestReaderWriter(unittest.TestCase): reader = Reader(MemorySource(buf.getvalue())) self.assertEqual(reader.header()["clustered"], False) + + def test_all_tiles(self): + buf = BytesIO() + writer = Writer(buf) + writer.write_tile(zxy_to_tileid(0, 0, 0), b"1") + writer.write_tile(zxy_to_tileid(1, 0, 0), b"1") + writer.write_tile(zxy_to_tileid(2, 0, 0), b"2") + writer.finalize( + { + "tile_compression": Compression.UNKNOWN, + "tile_type": TileType.UNKNOWN, + }, + {"key": "value"}, + ) + + reader = Reader(MemorySource(buf.getvalue())) + tiles = list(all_tiles(reader.get_bytes)) + self.assertEqual(tiles, [ + ((0,0,0), b"1"), + ((1,0,0), b"1"), + ((2,0,0), b"2"), + ]) \ No newline at end of file