mirrors/PMTiles
1
0
Fork 0
mirror of https://github.com/protomaps/PMTiles.git synced 2026-02-04 02:41:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor python reader to take a get_bytes function; add reader tests [#38]

Browse Source
This commit is contained in:
Brandon Liu
2022-04-26 16:15:03 +08:00
parent adb6f896a3
commit 3944d27833
4 changed files with 79 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
python/pmtiles/reader.py
View File

@@ -3,20 +3,26 @@ import mmap
from contextlib import contextmanager
@contextmanager
def read(fname):
r = Reader(fname)
try:
yield r
finally:
r.close()
def MmapSource(f):
mapping = mmap.mmap(f.fileno(), 0)
def get_bytes(offset, length):
return mapping[offset : offset + length]
return get_bytes
def MemorySource(buf):
def get_bytes(offset, length):
return buf[offset : offset + length]
return get_bytes
class Reader:
def __init__(self, fname):
self.f = open(fname, "r+b")
self.mmap = mmap.mmap(self.f.fileno(), 0)
assert int.from_bytes(self.mmap[0:2], byteorder="little") == 0x4D50
def __init__(self, get_bytes):
self.get_bytes = get_bytes
assert int.from_bytes(self.get_bytes(0, 2), byteorder="little") == 0x4D50
first_entry_idx = 10 + self.metadata_len
self.root_dir, self.leaves = self.load_directory(
first_entry_idx, self.root_entries
@@ -26,11 +32,11 @@ class Reader:
directory = {}
leaves = {}
for i in range(offset, offset + num_entries * 17, 17):
z = int.from_bytes(self.mmap[i : i + 1], byteorder="little")
x = int.from_bytes(self.mmap[i + 1 : i + 4], byteorder="little")
y = int.from_bytes(self.mmap[i + 4 : i + 7], byteorder="little")
tile_off = int.from_bytes(self.mmap[i + 7 : i + 13], byteorder="little")
tile_len = int.from_bytes(self.mmap[i + 13 : i + 17], byteorder="little")
z = int.from_bytes(self.get_bytes(i, 1), byteorder="little")
x = int.from_bytes(self.get_bytes(i + 1, 3), byteorder="little")
y = int.from_bytes(self.get_bytes(i + 4, 3), byteorder="little")
tile_off = int.from_bytes(self.get_bytes(i + 7, 6), byteorder="little")
tile_len = int.from_bytes(self.get_bytes(i + 13, 4), byteorder="little")
if z & 0b10000000:
leaves[(z & 0b01111111, x, y)] = (tile_off, tile_len)
else:
@@ -42,20 +48,20 @@ class Reader:
@property
def metadata_len(self):
return int.from_bytes(self.mmap[4:8], byteorder="little")
return int.from_bytes(self.get_bytes(4, 4), byteorder="little")
@property
def metadata(self):
s = self.mmap[10 : 10 + self.metadata_len]
s = self.get_bytes(10, self.metadata_len)
return json.loads(s)
@property
def version(self):
return int.from_bytes(self.mmap[2:4], byteorder="little")
return int.from_bytes(self.get_bytes(2, 2), byteorder="little")
@property
def root_entries(self):
return int.from_bytes(self.mmap[8:10], byteorder="little")
return int.from_bytes(self.get_bytes(8, 2), byteorder="little")
@property
def leaf_level(self):
@@ -64,7 +70,7 @@ class Reader:
def get(self, z, x, y):
val = self.root_dir.get((z, x, y))
if val:
return self.mmap[val[0] : val[0] + val[1]]
return self.get_bytes(val[0], val[1])
else:
if len(self.leaves) > 0:
level_diff = z - self.leaf_level
@@ -74,12 +80,12 @@ class Reader:
directory, _ = self.load_directory(val[0], val[1] // 17)
val = directory.get((z, x, y))
if val:
return self.mmap[val[0] : val[0] + val[1]]
return self.get_bytes(val[0], val[1])
def tiles(self):
for k, v in self.root_dir.items():
yield (k, self.mmap[v[0] : v[0] + v[1]])
yield (k, self.get_bytes(v[0], v[1]))
for val in self.leaves.values():
leaf_dir, _ = self.load_directory(val[0], val[1] // 17)
for k, v in leaf_dir.items():
yield (k, self.mmap[v[0] : v[0] + v[1]])
yield (k, self.get_bytes(v[0], v[1]))

27
python/pmtiles/writer.py
View File

@@ -80,20 +80,22 @@ def make_pyramid(tile_entries, start_leaf_offset, max_dir_size=21845):
@contextmanager
def write(fname):
w = Writer(fname)
f = open(fname, "wb")
w = Writer(f)
try:
yield w
finally:
w.close()
f.close()
class Writer:
def __init__(self, fname):
self.f = open(fname, "wb")
def __init__(self, f, max_dir_size):
self.offset = 512000
self.f = f
self.f.write(b"\0" * self.offset)
self.tile_entries = []
self.hash_to_offset = {}
self.max_dir_size = max_dir_size
def write_tile(self, z, x, y, data):
hsh = hash(data)
@@ -107,7 +109,7 @@ class Writer:
self.hash_to_offset[hsh] = self.offset
self.offset = self.offset + len(data)
def write_entry(self, entry):
def _write_entry(self, entry):
if entry.is_dir:
z_bytes = 0b10000000 | entry.z
else:
@@ -118,7 +120,7 @@ class Writer:
self.f.write(entry.offset.to_bytes(6, byteorder="little"))
self.f.write(entry.length.to_bytes(4, byteorder="little"))
def write_header(self, metadata, root_entries_len):
def _write_header(self, metadata, root_entries_len):
self.f.write((0x4D50).to_bytes(2, byteorder="little"))
self.f.write((2).to_bytes(2, byteorder="little"))
metadata_serialized = json.dumps(metadata)
@@ -129,24 +131,23 @@ class Writer:
self.f.write(metadata_serialized.encode("utf-8"))
def finalize(self, metadata={}):
root_dir, leaf_dirs = make_pyramid(self.tile_entries, self.offset)
root_dir, leaf_dirs = make_pyramid(
self.tile_entries, self.offset, self.max_dir_size
)
if len(leaf_dirs) > 0:
for leaf_dir in leaf_dirs:
for entry in leaf_dir:
self.write_entry(entry)
self._write_entry(entry)
self.f.seek(0)
self.write_header(metadata, len(root_dir))
self._write_header(metadata, len(root_dir))
for entry in root_dir:
self.write_entry(entry)
self._write_entry(entry)
return {
"num_tiles": len(self.tile_entries),
"num_unique_tiles": len(self.hash_to_offset),
"num_leaves": len(leaf_dirs),
}
def close(self):
self.f.close()