allow SDObject and Unit to be pickled

Author: aleivag

pystemd/base.py

@@ -27,6 +27,17 @@ def __init__(self, destination, path, bus=None, _autoload=False):
         if _autoload:
             self.load()
 
+    def __getstate__(self):
+        return {
+            "destination": self.destination,
+            "path": self.path,
+            "bus": self._bus,
+            "_autoload": self._loaded,
+        }
+
+    def __setstate__(self, state):
+        self.__init__(**state)
+
     def __enter__(self):
         self.load()
         return self
@@ -112,6 +123,7 @@ def load(self, force=False):
                 )
             elif interface_name == "org.freedesktop.DBus.Properties":
                 self.Properties = self._interfaces[interface_name]
+        self._loaded = True
 
 
 class SDInterface(object):

pystemd/machine1/machine.py

@@ -22,3 +22,10 @@ def __init__(self, external_id, bus=None, _autoload=False):
             bus=bus,
             _autoload=_autoload,
         )
+
+    def __getstate__(self):
+        return {
+            "external_id": self.external_id,
+            "bus": self._bus,
+            "_autoload": self._loaded,
+        }

pystemd/systemd1/unit.py

@@ -21,3 +21,10 @@ def __init__(self, external_id, bus=None, _autoload=False):
             bus=bus,
             _autoload=_autoload,
         )
+
+    def __getstate__(self):
+        return {
+            "external_id": self.external_id,
+            "bus": self._bus,
+            "_autoload": self._loaded,
+        }

tests/test_pickle.py

@@ -0,0 +1,19 @@
+import pickle
+
+from pystemd.systemd1.unit import Unit
+
+
+def test_unloaded_unit():
+    unit = Unit("foo.service")
+    sour_unit = pickle.loads(pickle.dumps(unit))
+    assert unit.external_id == sour_unit.external_id
+    assert not sour_unit._loaded
+
+
+def test_loaded_unit():
+    unit = Unit("foo.service")
+    unit.load()
+    sour_unit = pickle.loads(pickle.dumps(unit))
+    assert unit.external_id == sour_unit.external_id
+    assert sour_unit._loaded
+    assert sour_unit._interfaces

tests/test_version.py

@@ -1,6 +1,6 @@
 from pathlib import Path
-import toml
 
+import toml
 from cstq import Query