test_pickle.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  1. import pathlib
  2. import pickle
  3. import warnings
  4. from pickle import HIGHEST_PROTOCOL, dumps, loads
  5. import pytest
  6. import shapely
  7. from shapely import wkt
  8. from shapely.geometry import (
  9. GeometryCollection,
  10. LinearRing,
  11. LineString,
  12. MultiLineString,
  13. MultiPoint,
  14. MultiPolygon,
  15. Point,
  16. Polygon,
  17. box,
  18. )
  19. HERE = pathlib.Path(__file__).parent
  20. TEST_DATA = {
  21. "point2d": Point([(1.0, 2.0)]),
  22. "point3d": Point([(1.0, 2.0, 3.0)]),
  23. "linestring": LineString([(0.0, 0.0), (0.0, 1.0), (1.0, 1.0)]),
  24. "linearring": LinearRing([(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 0.0)]),
  25. "polygon": Polygon([(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 0.0)]),
  26. "multipoint": MultiPoint([(1.0, 2.0), (3.0, 4.0), (5.0, 6.0)]),
  27. "multilinestring": MultiLineString(
  28. [[(0.0, 0.0), (1.0, 1.0)], [(1.0, 2.0), (3.0, 3.0)]]
  29. ),
  30. "multipolygon": MultiPolygon([box(0, 0, 1, 1), box(2, 2, 3, 3)]),
  31. "geometrycollection": GeometryCollection([Point(1.0, 2.0), box(0, 0, 1, 1)]),
  32. "emptypoint": wkt.loads("POINT EMPTY"),
  33. "emptypolygon": wkt.loads("POLYGON EMPTY"),
  34. }
  35. TEST_NAMES, TEST_GEOMS = zip(*TEST_DATA.items())
  36. @pytest.mark.parametrize("geom1", TEST_GEOMS, ids=TEST_NAMES)
  37. def test_pickle_round_trip(geom1):
  38. data = dumps(geom1, HIGHEST_PROTOCOL)
  39. with warnings.catch_warnings():
  40. warnings.simplefilter("error")
  41. geom2 = loads(data)
  42. assert geom2.has_z == geom1.has_z
  43. assert type(geom2) is type(geom1)
  44. assert geom2.geom_type == geom1.geom_type
  45. assert geom2.wkt == geom1.wkt
  46. @pytest.mark.parametrize(
  47. "fname", (HERE / "data").glob("*.pickle"), ids=lambda fname: fname.name
  48. )
  49. def test_unpickle_pre_20(fname):
  50. from shapely.testing import assert_geometries_equal
  51. geom_type = fname.name.split("_")[0]
  52. expected = TEST_DATA[geom_type]
  53. with open(fname, "rb") as f:
  54. with pytest.warns(UserWarning, match="may be removed in a future version"):
  55. result = pickle.load(f)
  56. assert_geometries_equal(result, expected)
  57. if __name__ == "__main__":
  58. datadir = HERE / "data"
  59. datadir.mkdir(exist_ok=True)
  60. shapely_version = shapely.__version__
  61. print(shapely_version)
  62. print(shapely.geos_version)
  63. for name, geom in TEST_DATA.items():
  64. with open(datadir / f"{name}_{shapely_version}.pickle", "wb") as f:
  65. pickle.dump(geom, f)