Add to dict method to manage Enum serialization

Author: dgarros

backend/infrahub/core/schema/attribute_schema.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import enum
+from enum import Enum
 from typing import TYPE_CHECKING, Any, Optional, Union
 
-from pydantic import ConfigDict, field_validator, model_validator
+from pydantic import field_validator, model_validator
 
 from infrahub import config
 from infrahub.core.enums import generate_python_enum
@@ -21,8 +22,6 @@
 
 
 class AttributeSchema(GeneratedAttributeSchema):
-    model_config = ConfigDict(use_enum_values=True)
-
     _sort_by: list[str] = ["name"]
     _enum_class: Optional[type[enum.Enum]] = None
 
@@ -38,6 +37,13 @@ def is_relationship(self) -> bool:
     def is_deprecated(self) -> bool:
         return bool(self.deprecation)
 
+    def to_dict(self) -> dict:
+        data = self.model_dump(exclude_unset=True, exclude_none=True, exclude_defaults=True)
+        for field_name, value in data.items():
+            if isinstance(value, Enum):
+                data[field_name] = value.value
+        return data
+
     @field_validator("kind")
     @classmethod
     def kind_options(cls, v: str) -> str:

backend/infrahub/core/schema/basenode_schema.py

@@ -4,10 +4,11 @@
 import keyword
 import os
 from dataclasses import asdict, dataclass
+from enum import Enum
 from typing import TYPE_CHECKING, Any, Callable, Iterable, Literal, Optional, Union, overload
 
 from infrahub_sdk.utils import compare_lists, intersection
-from pydantic import ConfigDict, field_validator
+from pydantic import field_validator
 
 from infrahub.core.constants import RelationshipCardinality, RelationshipKind
 from infrahub.core.models import HashableModelDiff
@@ -27,8 +28,6 @@
 
 
 class BaseNodeSchema(GeneratedBaseNodeSchema):
-    model_config = ConfigDict(use_enum_values=True)
-
     _exclude_from_hash: list[str] = ["attributes", "relationships"]
     _sort_by: list[str] = ["namespace", "name"]
 
@@ -65,7 +64,15 @@ def __hash__(self) -> int:
         return hash(self.get_hash())
 
     def to_dict(self) -> dict:
-        return self.model_dump(exclude_unset=True, exclude_none=True, exclude_defaults=True)
+        data = self.model_dump(
+            exclude_unset=True, exclude_none=True, exclude_defaults=True, exclude={"attributes", "relationships"}
+        )
+        for field_name, value in data.items():
+            if isinstance(value, Enum):
+                data[field_name] = value.value
+        data["attributes"] = [attr.to_dict() for attr in self.attributes]
+        data["relationships"] = [rel.to_dict() for rel in self.relationships]
+        return data
 
     def get_hash(self, display_values: bool = False) -> str:
         """Extend the Hash Calculation to account for attributes and relationships."""

backend/infrahub/core/schema/relationship_schema.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+from enum import Enum
 from typing import TYPE_CHECKING, Any, Optional, Union
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel
 
 from infrahub import config
 from infrahub.core.constants import RelationshipDirection
@@ -19,8 +20,6 @@
 
 
 class RelationshipSchema(GeneratedRelationshipSchema):
-    model_config = ConfigDict(use_enum_values=True)
-
     _exclude_from_hash: list[str] = ["filters"]
     _sort_by: list[str] = ["name"]
 
@@ -36,6 +35,13 @@ def is_relationship(self) -> bool:
     def is_deprecated(self) -> bool:
         return bool(self.deprecation)
 
+    def to_dict(self) -> dict:
+        data = self.model_dump(exclude_unset=True, exclude_none=True, exclude_defaults=True)
+        for field_name, value in data.items():
+            if isinstance(value, Enum):
+                data[field_name] = value.value
+        return data
+
     def get_class(self) -> type[Relationship]:
         return Relationship