pyshp/tests/run_benchmarks.py at master · GeospatialPython/pyshp · GitHub

History

135 lines (110 loc) · 4.28 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

# Based on Taneli Hukkinen's https://github.com/hukkin/tomli-w/blob/master/benchmark/run.py

from __future__ import annotations

import collections

import functools

import os

import timeit

from collections.abc import Callable, Iterable

from os import PathLike

from pathlib import Path

from tempfile import TemporaryFile as TempF

from typing import cast, Iterable

import shapefile

# For shapefiles from https://github.com/JamesParrott/PyShp_test_shapefile

DEFAULT_PYSHP_TEST_REPO = (

rf"{os.getenv('USERPROFILE')}\Coding\repos\PyShp_test_shapefile"

)

PYSHP_TEST_REPO = Path(os.getenv("PYSHP_TEST_REPO", DEFAULT_PYSHP_TEST_REPO))

REPO_ROOT = Path(__file__).parent

blockgroups_file = REPO_ROOT / "shapefiles" / "blockgroups.shp"

edit_file = REPO_ROOT / "shapefiles" / "edit.shp"

merge_file = REPO_ROOT / "shapefiles" / "merge.shp"

states_provinces_file = PYSHP_TEST_REPO / "ne_10m_admin_1_states_provinces.shp"

tiny_countries_file = PYSHP_TEST_REPO / "ne_110m_admin_0_tiny_countries.shp"

gis_osm_natural_file = PYSHP_TEST_REPO / "gis_osm_natural_a_free_1.zip"

def benchmark(

name: str,

run_count: int,

func: Callable,

col_widths: tuple,

compare_to: float | None = None,

) -> float:

placeholder = "Running..."

print(f"{name:>{col_widths[0]}} | {placeholder}", end="", flush=True)

time_taken = timeit.timeit(func, number=run_count)

print("\b" * len(placeholder), end="")

time_suffix = " s"

print(f"{time_taken:{col_widths[1] - len(time_suffix)}.3g}{time_suffix}", end="")

print()

return time_taken

fields: dict[str | PathLike, list[shapefile.Field]] = {}

shapeRecords: dict[str | PathLike, list[shapefile.shapeRecord]] = (

collections.defaultdict(list)

)

def open_shapefile_with_PyShp(target: str | PathLike):

with shapefile.Reader(target) as r:

fields[target] = r.data_fields

for shapeRecord in r.iterShapeRecords():

shapeRecords[target].append(shapeRecord)

def write_shapefile_with_PyShp(target: str | PathLike):

with TempF("wb") as shp, TempF("wb") as dbf, TempF("wb") as shx:

with shapefile.Writer(shp=shp, dbf=dbf, shx=shx, strict=False) as w: # type: ignore [arg-type]

for field_info_tuple in fields[target]:

w.field(*field_info_tuple)

for shapeRecord in shapeRecords[target]:

w.shape(cast(shapefile.Shape, shapeRecord.shape))

record = cast(Iterable, shapeRecord.record)

w.record(*record)

SHAPEFILES = {

"Blockgroups": blockgroups_file,

"Edit": edit_file,

"Merge": merge_file,

"States_35MB": states_provinces_file,

"Tiny Countries": tiny_countries_file,

"GIS_OSM_zip_10MB": gis_osm_natural_file,

}

# Load files to avoid one off delays that only affect first disk seek

for file_path in SHAPEFILES.values():

file_path.read_bytes()

COLS_WIDTHS = (22, 10)

reader_benchmarks = {

test_name : functools.partial(

benchmark,

name=f"Read {test_name}",

func=functools.partial(open_shapefile_with_PyShp, target=target),

col_widths=COLS_WIDTHS,

)

for test_name, target in SHAPEFILES.items()

}

# Require fields and shapeRecords to first have been populated

# from data from previouly running the reader_benchmarks

writer_benchmarks = {

test_name : functools.partial(

benchmark,

name=f"Write {test_name}",

func=functools.partial(write_shapefile_with_PyShp, target=target),

col_widths=COLS_WIDTHS,

)

for test_name, target in SHAPEFILES.items()

}

def run(

run_count: int,

benchmarks: Iterable[Callable[[], None]],

col_widths: tuple[int, int] = COLS_WIDTHS,

) -> None:

col_head = ("parser", "exec time", "performance (more is better)")

print(f"Running benchmarks {run_count} times:")

print("-" * col_widths[0] + "---" + "-" * col_widths[1])

print(f"{col_head[0]:>{col_widths[0]}} | {col_head[1]:>{col_widths[1]}}")

print("-" * col_widths[0] + "-+-" + "-" * col_widths[1])

for benchmark in benchmarks:

benchmark( # type: ignore [call-arg]

run_count=run_count,

)

def main():

print("Reader tests:")

run(1, reader_benchmarks.values()) # type: ignore [arg-type]

print("\n\nWriter tests:")

run(1, writer_benchmarks.values()) # type: ignore [arg-type]

if __name__ == "__main__":

main()

pyshp/tests/run_benchmarks.py at master · GeospatialPython/pyshp · GitHub

Uh oh!

Expand file tree

Breadcrumbs

run_benchmarks.py

Latest commit

History

Breadcrumbs

run_benchmarks.py

File metadata and controls