|
| 1 | +import os |
| 2 | + |
| 3 | +from loguru import logger |
| 4 | + |
| 5 | +from pybind11_geobuf import rapidjson # noqa |
| 6 | +from pybind11_geobuf import Decoder, Encoder # noqa |
| 7 | +from pybind11_geobuf import normalize_json as normalize_json_impl # noqa |
| 8 | +from pybind11_geobuf import pbf_decode as pbf_decode_impl # noqa |
| 9 | + |
| 10 | + |
| 11 | +def __filesize(path: str) -> int: |
| 12 | + return os.stat(path).st_size |
| 13 | + |
| 14 | + |
| 15 | +def geobuf2json( |
| 16 | + input_path: str, |
| 17 | + output_path: str, |
| 18 | + *, |
| 19 | + indent: bool = False, |
| 20 | + sort_keys: bool = False, |
| 21 | +): |
| 22 | + logger.info( |
| 23 | + f"geobuf decoding {input_path} ({__filesize(input_path):,} bytes)..." |
| 24 | + ) # noqa |
| 25 | + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) |
| 26 | + decoder = Decoder() |
| 27 | + assert decoder.decode( |
| 28 | + geobuf=input_path, |
| 29 | + geojson=output_path, |
| 30 | + indent=indent, |
| 31 | + sort_keys=sort_keys, |
| 32 | + ), f"failed at decoding geojson, input:{input_path}, output:{output_path}" |
| 33 | + logger.info(f"wrote to {output_path} ({__filesize(output_path):,} bytes)") |
| 34 | + |
| 35 | + |
| 36 | +def json2geobuf( |
| 37 | + input_path: str, |
| 38 | + output_path: str, |
| 39 | + *, |
| 40 | + precision: int = 8, |
| 41 | +): |
| 42 | + logger.info( |
| 43 | + f"geobuf encoding {input_path} ({__filesize(input_path):,} bytes)..." |
| 44 | + ) # noqa |
| 45 | + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) |
| 46 | + encoder = Encoder(max_precision=int(10**precision)) |
| 47 | + assert encoder.encode( |
| 48 | + geojson=input_path, |
| 49 | + geobuf=output_path, |
| 50 | + ), f"failed at encoding geojson, input:{input_path}, output:{output_path}" |
| 51 | + logger.info(f"wrote to {output_path} ({__filesize(output_path):,} bytes)") |
| 52 | + |
| 53 | + |
| 54 | +def normalize_geobuf( |
| 55 | + input_path: str, |
| 56 | + output_path: str = None, |
| 57 | + *, |
| 58 | + sort_keys: bool = True, |
| 59 | + precision: int = -1, |
| 60 | +): |
| 61 | + logger.info( |
| 62 | + f"normalize_geobuf {input_path} ({__filesize(input_path):,} bytes)" |
| 63 | + ) # noqa |
| 64 | + with open(input_path, "rb") as f: |
| 65 | + encoded = f.read() |
| 66 | + decoder = Decoder() |
| 67 | + json = decoder.decode_to_rapidjson(encoded, sort_keys=sort_keys) |
| 68 | + if precision < 0: |
| 69 | + precision = decoder.precision() |
| 70 | + logger.info(f"auto precision from geobuf: {precision}") |
| 71 | + else: |
| 72 | + logger.info(f"user precision: {precision}") |
| 73 | + encoder = Encoder(max_precision=int(10**precision)) |
| 74 | + encoded = encoder.encode(json) |
| 75 | + logger.info(f"encoded #bytes: {len(encoded):,}") |
| 76 | + output_path = output_path or input_path |
| 77 | + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) |
| 78 | + with open(output_path, "wb") as f: |
| 79 | + f.write(encoded) |
| 80 | + logger.info(f"wrote to {output_path} ({__filesize(output_path):,} bytes)") |
| 81 | + |
| 82 | + |
| 83 | +def normalize_json( |
| 84 | + input_path: str, |
| 85 | + output_path: str = None, |
| 86 | + *, |
| 87 | + indent: bool = True, |
| 88 | + sort_keys: bool = True, |
| 89 | + precision: int = -1, |
| 90 | +): |
| 91 | + logger.info( |
| 92 | + f"normalize_json {input_path} ({__filesize(input_path):,} bytes)" |
| 93 | + ) # noqa |
| 94 | + output_path = output_path or input_path |
| 95 | + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) |
| 96 | + if precision > 0: |
| 97 | + logger.info( |
| 98 | + f"convert to geobuf (precision: {precision}), then back to geojson" |
| 99 | + ) # noqa |
| 100 | + encoder = Encoder(max_precision=int(10**precision)) |
| 101 | + geojson = rapidjson().load(input_path) |
| 102 | + assert geojson.IsObject(), f"invalid geojson: {input_path}" |
| 103 | + encoded = encoder.encode(geojson) |
| 104 | + logger.info(f"geobuf #bytes: {len(encoded):,}") |
| 105 | + decoder = Decoder() |
| 106 | + text = decoder.decode(encoded, indent=indent, sort_keys=sort_keys) |
| 107 | + logger.info(f"geojson #bytes: {len(text):,}") |
| 108 | + with open(output_path, "w", encoding="utf8") as f: |
| 109 | + f.write(text) |
| 110 | + else: |
| 111 | + assert normalize_json_impl( |
| 112 | + input_path, |
| 113 | + output_path, |
| 114 | + indent=indent, |
| 115 | + sort_keys=sort_keys, |
| 116 | + ), f"failed to normalize json to {output_path}" |
| 117 | + logger.info(f"wrote to {output_path} ({__filesize(output_path):,} bytes)") |
| 118 | + |
| 119 | + |
| 120 | +def pbf_decode(path: str, output_path: str = None, *, indent: str = ""): |
| 121 | + with open(path, "rb") as f: |
| 122 | + data = f.read() |
| 123 | + decoded = pbf_decode_impl(data, indent=indent) |
| 124 | + if output_path: |
| 125 | + os.makedirs( |
| 126 | + os.path.dirname(os.path.abspath(output_path)), exist_ok=True |
| 127 | + ) # noqa |
| 128 | + with open(output_path, "w", encoding="utf8") as f: |
| 129 | + f.write(decoded) |
| 130 | + logger.info(f"wrote to {output_path}") |
| 131 | + else: |
| 132 | + print(decoded) |
| 133 | + |
| 134 | + |
| 135 | +if __name__ == "__main__": |
| 136 | + import fire |
| 137 | + |
| 138 | + fire.core.Display = lambda lines, out: print(*lines, file=out) |
| 139 | + fire.Fire( |
| 140 | + { |
| 141 | + "geobuf2json": geobuf2json, |
| 142 | + "json2geobuf": json2geobuf, |
| 143 | + "normalize_geobuf": normalize_geobuf, |
| 144 | + "normalize_json": normalize_json, |
| 145 | + "pbf_decode": pbf_decode, |
| 146 | + } |
| 147 | + ) |
0 commit comments