Make CStruct/MemCStruct Pickle Friendly

Author: andreax79

cstruct/abstract.py

@@ -229,3 +229,23 @@ def __str__(self) -> str:
 
     def __repr__(self) -> str:  # pragma: no cover
         return self.__str__()
+
+    def __getstate__(self) -> bytes:
+        """
+        This method is called and the returned object is pickled
+        as the contents for the instance, instead of the contents of
+        the instance’s dictionary
+
+        Returns:
+            bytes: The packed structure
+        """
+        return self.pack()
+
+    def __setstate__(self, state: bytes) -> bool:
+        """
+        This method it is called with the unpickled state
+
+        Args:
+            state: bytes to be unpacked
+        """
+        return self.unpack(state)

tests/test_pickle.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+
+import pickle
+import cstruct
+
+
+class Position(cstruct.CStruct):
+    __byte_order__ = cstruct.LITTLE_ENDIAN
+    __def__ = """
+        struct {
+            unsigned char head;
+            unsigned char sector;
+            unsigned char cyl;
+        }
+    """
+
+
+class MemPosition(cstruct.MemCStruct):
+    __byte_order__ = cstruct.LITTLE_ENDIAN
+    __def__ = """
+        struct {
+            unsigned char head;
+            unsigned char sector;
+            unsigned char cyl;
+        }
+    """
+
+
+def test_pickle():
+    # pickle an object
+    original_pos = Position(head=254, sector=63, cyl=134)
+    pickled_bytes = pickle.dumps(original_pos)
+
+    # reconstitute a pickled object
+    reconstituted_pos = pickle.loads(pickled_bytes)
+
+    assert reconstituted_pos.head == original_pos.head
+    assert reconstituted_pos.sector == original_pos.sector
+    assert reconstituted_pos.cyl == original_pos.cyl
+
+
+def test_mem_pickle():
+    # pickle an object
+    original_pos = MemPosition(head=254, sector=63, cyl=134)
+    pickled_bytes = pickle.dumps(original_pos)
+
+    # reconstitute a pickled object
+    reconstituted_pos = pickle.loads(pickled_bytes)
+
+    assert reconstituted_pos.head == original_pos.head
+    assert reconstituted_pos.sector == original_pos.sector
+    assert reconstituted_pos.cyl == original_pos.cyl