From 88886e3fd63a0c3c9d5c1577b913c7d1f8326406 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mose=20M=C3=BCller?= <mosemueller@gmail.com>
Date: Mon, 11 Dec 2023 17:25:03 +0100
Subject: [PATCH] fixed serialization of class deriving from class which
 derives from DataService

---
 src/pydase/utils/helpers.py    | 20 +++++++++++---------
 src/pydase/utils/serializer.py | 31 +++++++++++++++++--------------
 tests/utils/test_serializer.py | 24 ++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 23 deletions(-)

diff --git a/src/pydase/utils/helpers.py b/src/pydase/utils/helpers.py
index c4aa2be..92fb6ae 100644
--- a/src/pydase/utils/helpers.py
+++ b/src/pydase/utils/helpers.py
@@ -177,19 +177,21 @@ def parse_list_attr_and_index(attr_string: str) -> tuple[str, int | None]:
     return attr_name, index
 
 
-def get_component_class_names() -> list[str]:
+def get_component_classes() -> list[type]:
     """
-    Returns the names of the component classes in a list.
-
-    It takes the names from the pydase/components/__init__.py file, so this file should
-    always be up-to-date with the currently available components.
-
-    Returns:
-        list[str]: List of component class names
+    Returns references to the component classes in a list.
     """
     import pydase.components
 
-    return pydase.components.__all__
+    return [
+        getattr(pydase.components, cls_name) for cls_name in pydase.components.__all__
+    ]
+
+
+def get_data_service_class_reference() -> Any:
+    import pydase.data_service.data_service
+
+    return getattr(pydase.data_service.data_service, "DataService")
 
 
 def is_property_attribute(target_obj: Any, attr_name: str) -> bool:
diff --git a/src/pydase/utils/serializer.py b/src/pydase/utils/serializer.py
index 6513368..90a57ca 100644
--- a/src/pydase/utils/serializer.py
+++ b/src/pydase/utils/serializer.py
@@ -8,7 +8,8 @@ import pydase.units as u
 from pydase.data_service.abstract_data_service import AbstractDataService
 from pydase.utils.helpers import (
     get_attribute_doc,
-    get_component_class_names,
+    get_component_classes,
+    get_data_service_class_reference,
     parse_list_attr_and_index,
 )
 
@@ -157,24 +158,26 @@ class Serializer:
     def _serialize_data_service(obj: AbstractDataService) -> dict[str, Any]:
         readonly = False
         doc = get_attribute_doc(obj)
-        obj_type = type(obj).__name__
-        if type(obj).__name__ not in get_component_class_names():
-            obj_type = "DataService"
+        obj_type = "DataService"
 
-        # Get the dictionary of the base class
-        base_set = set(type(obj).__base__.__dict__)
-        # Get the dictionary of the derived class
-        derived_set = set(type(obj).__dict__)
-        # Get the difference between the two dictionaries
-        derived_only_set = derived_set - base_set
+        # Get component base class if any
+        component_base_cls = next(
+            (cls for cls in get_component_classes() if isinstance(obj, cls)), None
+        )
+        if component_base_cls:
+            obj_type = component_base_cls.__name__
+
+        # Get the set of DataService class attributes
+        data_service_attr_set = set(dir(get_data_service_class_reference()))
+        # Get the set of the object attributes
+        obj_attr_set = set(dir(obj))
+        # Get the difference between the two sets
+        derived_only_attr_set = obj_attr_set - data_service_attr_set
 
-        instance_dict = set(obj.__dict__)
-        # Merge the class and instance dictionaries
-        merged_set = derived_only_set | instance_dict
         value = {}
 
         # Iterate over attributes, properties, class attributes, and methods
-        for key in sorted(merged_set):
+        for key in sorted(derived_only_attr_set):
             if key.startswith("_"):
                 continue  # Skip attributes that start with underscore
 
diff --git a/tests/utils/test_serializer.py b/tests/utils/test_serializer.py
index 8db0ddb..797de9d 100644
--- a/tests/utils/test_serializer.py
+++ b/tests/utils/test_serializer.py
@@ -294,6 +294,30 @@ def test_dict_serialization() -> None:
     }
 
 
+def test_derived_data_service_serialization() -> None:
+    class BaseService(pydase.DataService):
+        class_attr = 1337
+
+        def __init__(self) -> None:
+            super().__init__()
+            self._name = "Service"
+
+        @property
+        def name(self) -> str:
+            return self._name
+
+        @name.setter
+        def name(self, value: str) -> None:
+            self._name = value
+
+    class DerivedService(BaseService):
+        ...
+
+    base_instance = BaseService()
+    service_instance = DerivedService()
+    assert service_instance.serialize() == base_instance.serialize()
+
+
 @pytest.fixture
 def setup_dict():
     class MySubclass(pydase.DataService):