|
| 1 | +# Based on Taneli Hukkinen's https://github.com/hukkin/tomli-w/blob/master/benchmark/run.py |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import functools |
| 6 | +import os |
| 7 | +import timeit |
| 8 | +from collections.abc import Callable |
| 9 | +from pathlib import Path |
| 10 | +from typing import Union |
| 11 | + |
| 12 | +import shapefile as shp |
| 13 | + |
| 14 | +# For shapefiles from https://github.com/JamesParrott/PyShp_test_shapefile |
| 15 | +DEFAULT_PYSHP_TEST_REPO = ( |
| 16 | + rf"{os.getenv('USERPROFILE')}\Coding\repos\PyShp_test_shapefile" |
| 17 | +) |
| 18 | +PYSHP_TEST_REPO = Path(os.getenv("PYSHP_TEST_REPO", DEFAULT_PYSHP_TEST_REPO)) |
| 19 | +REPO_ROOT = Path(__file__).parent |
| 20 | + |
| 21 | + |
| 22 | +blockgroups_file = REPO_ROOT / "shapefiles" / "blockgroups.shp" |
| 23 | +edit_file = REPO_ROOT / "shapefiles" / "test" / "edit.shp" |
| 24 | +merge_file = REPO_ROOT / "shapefiles" / "test" / "merge.shp" |
| 25 | +states_provinces_file = PYSHP_TEST_REPO / "ne_10m_admin_1_states_provinces.shp" |
| 26 | +tiny_countries_file = PYSHP_TEST_REPO / "ne_110m_admin_0_tiny_countries.shp" |
| 27 | +gis_osm_natural_file = PYSHP_TEST_REPO / "gis_osm_natural_a_free_1.zip" |
| 28 | + |
| 29 | + |
| 30 | +def benchmark( |
| 31 | + name: str, |
| 32 | + run_count: int, |
| 33 | + func: Callable, |
| 34 | + col_width: tuple, |
| 35 | + compare_to: float | None = None, |
| 36 | +) -> float: |
| 37 | + placeholder = "Running..." |
| 38 | + print(f"{name:>{col_width[0]}} | {placeholder}", end="", flush=True) |
| 39 | + time_taken = timeit.timeit(func, number=run_count) |
| 40 | + print("\b" * len(placeholder), end="") |
| 41 | + time_suffix = " s" |
| 42 | + print(f"{time_taken:{col_width[1]-len(time_suffix)}.3g}{time_suffix}", end="") |
| 43 | + print() |
| 44 | + return time_taken |
| 45 | + |
| 46 | + |
| 47 | +def open_shapefile_with_PyShp(target: Union[str, os.PathLike]): |
| 48 | + with shp.Reader(target) as r: |
| 49 | + for shapeRecord in r.iterShapeRecords(): |
| 50 | + pass |
| 51 | + |
| 52 | + |
| 53 | +READER_TESTS = { |
| 54 | + "Blockgroups": blockgroups_file, |
| 55 | + "Edit": edit_file, |
| 56 | + "Merge": merge_file, |
| 57 | + "States_35MB": states_provinces_file, |
| 58 | + "Tiny Countries": tiny_countries_file, |
| 59 | + "GIS_OSM_zip_10MB": gis_osm_natural_file, |
| 60 | +} |
| 61 | + |
| 62 | + |
| 63 | +def run(run_count: int) -> None: |
| 64 | + col_width = (21, 10) |
| 65 | + col_head = ("parser", "exec time", "performance (more is better)") |
| 66 | + # Load files to avoid one off delays that only affect first disk seek |
| 67 | + for file_path in READER_TESTS.values(): |
| 68 | + file_path.read_bytes() |
| 69 | + print(f"Running benchmarks {run_count} times:") |
| 70 | + print("-" * col_width[0] + "---" + "-" * col_width[1]) |
| 71 | + print(f"{col_head[0]:>{col_width[0]}} | {col_head[1]:>{col_width[1]}}") |
| 72 | + print("-" * col_width[0] + "-+-" + "-" * col_width[1]) |
| 73 | + for test_name, target in READER_TESTS.items(): |
| 74 | + benchmark( |
| 75 | + f"Read {test_name}", |
| 76 | + run_count, |
| 77 | + functools.partial(open_shapefile_with_PyShp, target=target), |
| 78 | + col_width, |
| 79 | + ) |
| 80 | + |
| 81 | + |
| 82 | +if __name__ == "__main__": |
| 83 | + run(1) |
0 commit comments