Make item pickle smaller (#1285)

* Make item pickle smaller

* Don't make all the links pickle in the same way

* Update changelog

CHANGELOG.md

@@ -2,9 +2,13 @@
 
 ## [Unreleased]
 
+### Changed
+
+- Made item pickles smaller by changing how nested links are stored([#1285](https://github.com/stac-utils/pystac/pull/1285))
+
 ### Fixed
 
- - No longer use the `datetime.utcnow` method that has been deprecated in Python 3.12 ([#1283](https://github.com/stac-utils/pystac/pull/1283))
+- No longer use the `datetime.utcnow` method that has been deprecated in Python 3.12 ([#1283](https://github.com/stac-utils/pystac/pull/1283))
 
 ## [v1.9.0] - 2023-10-23
 

pystac/item.py

@@ -172,6 +172,27 @@ def __init__(
     def __repr__(self) -> str:
         return f"<Item id={self.id}>"
 
+    def __getstate__(self) -> dict[str, Any]:
+        """Ensure that pystac does not encode too much information when pickling"""
+        d = self.__dict__.copy()
+
+        d["links"] = [
+            link.to_dict() if link.get_href() else link for link in d["links"]
+        ]
+
+        return d
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Ensure that pystac knows how to decode the pickled object"""
+        d = state.copy()
+
+        d["links"] = [
+            Link.from_dict(link).set_owner(self) if isinstance(link, dict) else link
+            for link in d["links"]
+        ]
+
+        self.__dict__ = d
+
     def set_self_href(self, href: str | None) -> None:
         """Sets the absolute HREF that is represented by the ``rel == 'self'``
         :class:`~pystac.Link`.

tests/test_item.py

@@ -2,6 +2,7 @@
 
 import json
 import os
+import pickle
 import tempfile
 import unittest
 from copy import deepcopy
@@ -644,3 +645,40 @@ def test_invalid_error_message(item: Item) -> None:
     with pytest.raises(STACValidationError) as error:
         item.validate()
     assert "can't have a collection" in str(error.value)
+
+
+def test_pickle_with_no_links(item: Item) -> None:
+    roundtripped = pickle.loads(pickle.dumps(item))
+    for attr in ["id", "geometry", "bbox", "datetime", "links"]:
+        assert getattr(roundtripped, attr) == getattr(item, attr)
+
+
+def test_pickle_with_hrefless_links(item: Item) -> None:
+    root = pystac.Catalog("root", "root")
+    a = pystac.Catalog("a", "a")
+
+    item.add_link(pystac.Link("related", a))
+    item.add_link(
+        pystac.Link("item", TestCases.get_path("data-files/item/sample-item.json"))
+    )
+    item.set_root(root)
+
+    roundtripped = pickle.loads(pickle.dumps(item))
+    for original, new in zip(item.links, roundtripped.links):
+        assert original.rel == new.rel
+        assert original.media_type == new.media_type
+        assert str(original.owner) == str(new.owner)
+        assert str(original.target) == str(new.target)
+
+
+def test_pickle_with_only_href_links(item: Item) -> None:
+    item.add_link(
+        pystac.Link("item", TestCases.get_path("data-files/item/sample-item.json"))
+    )
+
+    roundtripped = pickle.loads(pickle.dumps(item))
+    for original, new in zip(item.links, roundtripped.links):
+        assert original.rel == new.rel
+        assert original.media_type == new.media_type
+        assert str(original.owner) == str(new.owner)
+        assert str(original.target) == str(new.target)