From 3e557f08ce64f3c774715c6ff3721849c2f3b277 Mon Sep 17 00:00:00 2001
From: Alvaro Leiva <aleivag@gmail.com>
Date: Tue, 16 Sep 2025 09:35:50 -0700
Subject: [PATCH] allow SDObject and Unit to be pickled

---
 pystemd/base.py             | 12 ++++++++++++
 pystemd/machine1/machine.py |  7 +++++++
 pystemd/systemd1/unit.py    |  7 +++++++
 tests/test_pickle.py        | 19 +++++++++++++++++++
 tests/test_version.py       |  2 +-
 5 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_pickle.py

diff --git a/pystemd/base.py b/pystemd/base.py
index 2d59c00..03329a9 100644
--- a/pystemd/base.py
+++ b/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):
diff --git a/pystemd/machine1/machine.py b/pystemd/machine1/machine.py
index 20d2bf5..5f1abf6 100644
--- a/pystemd/machine1/machine.py
+++ b/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,
+        }
diff --git a/pystemd/systemd1/unit.py b/pystemd/systemd1/unit.py
index 4657adf..ff56fc9 100644
--- a/pystemd/systemd1/unit.py
+++ b/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,
+        }
diff --git a/tests/test_pickle.py b/tests/test_pickle.py
new file mode 100644
index 0000000..142a417
--- /dev/null
+++ b/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
diff --git a/tests/test_version.py b/tests/test_version.py
index 8b5d0d6..0afa52c 100644
--- a/tests/test_version.py
+++ b/tests/test_version.py
@@ -1,6 +1,6 @@
 from pathlib import Path
-import toml
 
+import toml
 from cstq import Query