From a981bff4179a1b8bdebf33f39e9a4b2096f210ae Mon Sep 17 00:00:00 2001
From: Rafael Leira <rafael.leira@naudit.es>
Date: Fri, 19 Apr 2024 16:39:51 +0200
Subject: [PATCH] Restored the getstate method on device.py. Requested in #86

---
 CHANGELOG.md                      |  3 ++-
 pySMART/device.py                 | 31 ++++++++++++++-----------------
 pySMART/interface/ata/__init__.py |  2 +-
 3 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2dcc976..7f8a484 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,7 +1,8 @@
 Version 1.3.1
 =============
 - Added support for test polling_minutes. Requested in [#77](https://github.com/truenas/py-SMART/issues/77) and [#78](https://github.com/truenas/py-SMART/pull/78). Thanks @Gigahawk
-- Minnor typo fixed in NvmeAttributes (issue [#81](https://github.com/truenas/py-SMART/issues/81)).  Thanks @petersulyok
+- Minnor typo fixed in NvmeAttributes (issue [#81](https://github.com/truenas/py-SMART/issues/81)). Thanks @petersulyok
+- Fixed __getstate__ method in Device class. (issue [#86](https://github.com/truenas/py-SMART/issues/86)). Thanks @f18m
 - **Breaking changes**
     - **smartctl.all**:
         - Officially removed support for python 3.7.
diff --git a/pySMART/device.py b/pySMART/device.py
index 2a508cf..29b2d44 100644
--- a/pySMART/device.py
+++ b/pySMART/device.py
@@ -481,36 +481,33 @@ def __repr__(self):
             self.serial
         )
 
-    def __getstate__(self, all_info=True):
+    def __getstate__(self, all_info=True) -> Dict:
         """
         Allows us to send a pySMART Device object over a serializable
         medium which uses json (or the likes of json) payloads
         """
-        return None
+
         state_dict = {
+            'attributes': [attr.__getstate__() if attr else None for attr in self.attributes],
+            'capacity': self._capacity_human,
+            'diagnostics': self.diagnostics.__getstate__(all_info),
+            'firmware': self.firmware,
+            'if_attributes': self.if_attributes.__getstate__() if self.if_attributes else None,
             'interface': self._interface if self._interface else 'UNKNOWN INTERFACE',
+            'is_ssd': self.is_ssd,
+            'messages': self.messages,
             'model': self.model,
-            'firmware': self.firmware,
+            'name': self.name,
+            'path': self.dev_reference,
+            'rotation_rate': self.rotation_rate,
+            'serial': self.serial,
             'smart_capable': self.smart_capable,
             'smart_enabled': self.smart_enabled,
             'smart_status': self.assessment,
-            'messages': self.messages,
+            'temperature': self.temperature,
             'test_capabilities': self.test_capabilities.copy(),
             'tests': [t.__getstate__() for t in self.tests] if self.tests else [],
-            'diagnostics': self.diagnostics.__getstate__(all_info),
-            'temperature': self.temperature,
-            'attributes': [attr.__getstate__() if attr else None for attr in self.attributes],
-            'if_attributes': self.if_attributes.__getstate__(all_info) if self.if_attributes else None,
         }
-        if all_info:
-            state_dict.update({
-                'name': self.name,
-                'path': self.dev_reference,
-                'serial': self.serial,
-                'is_ssd': self.is_ssd,
-                'rotation_rate': self.rotation_rate,
-                'capacity': self._capacity_human
-            })
         return state_dict
 
     def __setstate__(self, state):
diff --git a/pySMART/interface/ata/__init__.py b/pySMART/interface/ata/__init__.py
index c5b91b7..290115a 100644
--- a/pySMART/interface/ata/__init__.py
+++ b/pySMART/interface/ata/__init__.py
@@ -64,7 +64,7 @@ def parse(self, data: Iterator[str]) -> None:
                         self.legacyAttributes[int(tmp['id'])] = Attribute(
                             int(tmp['id']), tmp['name'], int(tmp['flag'], base=16), tmp['value'], tmp['worst'], tmp['thresh'], tmp['type'], tmp['updated'], tmp['whenfailed'], tmp['raw'])
 
-    def __getstate__(self, all_info=True):
+    def __getstate__(self):
         """
         Allows us to send a pySMART diagnostics object over a serializable
         medium which uses json (or the likes of json) payloads