Repository: benbovy/spherely (language: cpp)

Issue:
Support pickling Geography objects
At the moment we didn't implement this, so you will get `TypeError: cannot pickle 'spherely.Point' object`, but it would be useful to support it.

To provide a full exact roundtrip through pickling, we should maybe pickle the raw unit vector (x, y, z) of the `S2Point`, instead of going through latitude/longitude like other conversions (eg WKB), because there can always be some floating point difference in this conversion.




Discussion / hints:
S2geometry provides `Encoder` and `Decoder` objects as well as `Encode()` and `Decode()` methods from most of its data types (point, polyline, polygon) so we might be able to just reuse that. For a point it serializes the raw unit vector. It is also possible to encode `S2ShapeIndex` objects with lazy decoding.

Probably cleanest would be to expose this API via `s2geography::Geography` wrapper classes rather than dealing with it directly in spherely (same for other methods like `Clone()`, which would be better than [this function](https://github.com/benbovy/spherely/blob/165f0ef26fdc67126c83ef7996055636e3e5a3fc/src/geography.cpp#L35-L77) implemented in #51).
Actually, I think we could already access `Encode()` and `Decode()` via `s2geography::Geography::Shape()`.
https://github.com/paleolimbot/s2geography/pull/40 is now merged so adding support for pickle should be easier.

That s2geography PR also added the `GeographyKind` enum tag, which will allow us to clean up some parts of the code in Spherely (e.g., remove some `dynamic_cast`).



Fix patch:
diff --git a/src/geography.cpp b/src/geography.cpp
index d84c177..bd527fc 100644
--- a/src/geography.cpp
+++ b/src/geography.cpp
@@ -7,6 +7,7 @@
 #include <s2/s2loop.h>
 #include <s2/s2point.h>
 #include <s2/s2polygon.h>
+#include <s2/util/coding/coder.h>
 #include <s2geography/geography.h>
 #include <s2geography/predicates.h>
 #include <s2geography/wkt-writer.h>
@@ -171,6 +172,53 @@ void Geography::extract_geog_properties() {
     }
 }
 
+py::tuple Geography::encode() const {
+    // encode geography type
+    using IntType = std::underlying_type_t<GeographyType>;
+    auto encoded_geog_type = static_cast<IntType>(geog_type());
+
+    // encode empty
+    // (note: this is already handled internally by s2geography::Geography::EncodeTagged() but
+    // there no current way to get the that information externally when/after decoding)
+    auto empty = m_is_empty;
+
+    // encode geog
+    Encoder geog_encoder;
+    s2geog::EncodeOptions encode_opts;
+    geog().EncodeTagged(&geog_encoder, encode_opts);
+
+    std::string encoded_geog;
+    encoded_geog.assign(geog_encoder.base(), geog_encoder.base() + geog_encoder.length());
+
+    return py::make_tuple(encoded_geog_type, empty, py::bytes(encoded_geog));
+}
+
+Geography Geography::decode(const py::tuple &encoded) {
+    auto decoded = Geography();
+
+    // decode geography type
+    using IntType = std::underlying_type_t<GeographyType>;
+    GeographyType geog_type{encoded[0].cast<IntType>()};
+    decoded.m_geog_type = geog_type;
+
+    // decode empty
+    decoded.m_is_empty = encoded[1].cast<bool>();
+
+    // decode geog() (s2geography::Geography)
+    auto encoded_geog = encoded[2].cast<std::string>();
+    Decoder geog_decoder(encoded_geog.c_str(), encoded_geog.size());
+    auto decoded_geog_ptr = s2geog::Geography::DecodeTagged(&geog_decoder);
+
+    // TODO: remove this quick & dirty fix (https://github.com/paleolimbot/s2geography/issues/54)
+    if (decoded_geog_ptr->kind() == s2geog::GeographyKind::GEOGRAPHY_COLLECTION) {
+        decoded.m_s2geog_ptr = clone_s2geography(*decoded_geog_ptr);
+    } else {
+        decoded.m_s2geog_ptr = std::move(decoded_geog_ptr);
+    }
+
+    return decoded;
+}
+
 /*
 ** Geography properties
 */
@@ -267,6 +315,9 @@ void init_geography(py::module &m) {
         return s2geog::s2_equals(idx1, idx2, options);
     });
 
+    pygeography.def(py::pickle([](Geography &geog) { return geog.encode(); },
+                               [](py::tuple &encoded) { return Geography::decode(encoded); }));
+
     // Geography properties
 
     m.def("get_type_id",
diff --git a/src/geography.hpp b/src/geography.hpp
index bc116ca..c703647 100644
--- a/src/geography.hpp
+++ b/src/geography.hpp
@@ -107,6 +107,9 @@ class Geography {
     Geography clone() const;
     std::unique_ptr<s2geog::Geography> clone_geog() const;
 
+    py::tuple encode() const;
+    static Geography decode(const py::tuple& encoded);
+
 private:
     S2GeographyPtr m_s2geog_ptr;
     S2GeographyIndexPtr m_s2geog_index_ptr;


Test patch:
diff --git a/tests/test_geography.py b/tests/test_geography.py
index c812926..6713141 100644
--- a/tests/test_geography.py
+++ b/tests/test_geography.py
@@ -1,3 +1,5 @@
+import pickle
+
 import pytest
 import numpy as np
 
@@ -130,3 +132,37 @@ def test_equality() -> None:
     assert poly1 == poly2
     assert poly2 == poly3
     assert poly1 == poly3
+
+    coll1 = (spherely.create_collection([spherely.create_point(40, 50)]),)
+    coll2 = (spherely.create_collection([spherely.create_point(40, 50)]),)
+
+    assert coll1 == coll2
+
+
+@pytest.mark.parametrize(
+    "geog",
+    [
+        spherely.create_point(45, 50),
+        spherely.create_multipoint([(5, 50), (6, 51)]),
+        spherely.create_linestring([(5, 50), (6, 51)]),
+        spherely.create_multilinestring([[(5, 50), (6, 51)], [(15, 60), (16, 61)]]),
+        spherely.create_polygon([(5, 50), (5, 60), (6, 60), (6, 51)]),
+        spherely.create_multipolygon(
+            [
+                spherely.create_polygon([(5, 50), (5, 60), (6, 60), (6, 51)]),
+                spherely.create_polygon([(10, 100), (10, 160), (11, 160), (11, 100)]),
+            ]
+        ),
+        spherely.create_collection([spherely.create_point(40, 50)]),
+        # empty geography
+        spherely.create_point(),
+        spherely.create_linestring(),
+        spherely.create_polygon(),
+    ],
+)
+def test_pickle_roundtrip(geog):
+    roundtripped = pickle.loads(pickle.dumps(geog))
+
+    assert spherely.get_type_id(roundtripped) == spherely.get_type_id(geog)
+    assert spherely.to_wkt(roundtripped) == spherely.to_wkt(geog)
+    assert roundtripped == geog