Rename decorator to register_severity_calculator

5 files changed, 34 insertions, 28 deletions

diff --git a/examples/custom_severity_calculator.py b/examples/custom_severity_calculator.py
index 9754dd8..5978bfc 100644
--- a/examples/custom_severity_calculator.py
+++ b/examples/custom_severity_calculator.py
@@ -5,10 +5,10 @@ from __future__ import annotations
 from typing import Optional
 
 import robusta_krr
-from robusta_krr.api.models import Severity, ResourceType
+from robusta_krr.api.models import Severity, ResourceType, register_severity_calculator
 
 
-@Severity.bind_calculator(ResourceType.CPU)
+@register_severity_calculator(ResourceType.CPU)
 def percentage_severity_calculator(
     current: Optional[float], recommended: Optional[float], resource_type: ResourceType
 ) -> Severity:
diff --git a/robusta_krr/api/models.py b/robusta_krr/api/models.py
index c6b3c99..5340f29 100644
--- a/robusta_krr/api/models.py
+++ b/robusta_krr/api/models.py
@@ -2,7 +2,7 @@ from robusta_krr.core.abstract.strategies import HistoryData, ResourceRecommenda
 from robusta_krr.core.models.allocations import RecommendationValue, ResourceAllocations, ResourceType
 from robusta_krr.core.models.objects import K8sObjectData
 from robusta_krr.core.models.result import ResourceScan, Result
-from robusta_krr.core.models.severity import Severity
+from robusta_krr.core.models.severity import Severity, register_severity_calculator
 
 __all__ = [
     "ResourceType",
@@ -11,6 +11,7 @@ __all__ = [
     "K8sObjectData",
     "Result",
     "Severity",
+    "register_severity_calculator",
     "bind_calculator",
     "ResourceScan",
     "ResourceRecommendation",
diff --git a/robusta_krr/core/models/objects.py b/robusta_krr/core/models/objects.py
index 21f5d61..8948319 100644
--- a/robusta_krr/core/models/objects.py
+++ b/robusta_krr/core/models/objects.py
@@ -39,3 +39,8 @@ class K8sObjectData(pd.BaseModel):
     @property
     def pods_count(self) -> int:
         return len(self.pods)
+
+    class Config:
+        json_encoders = {
+            PodData: lambda pod: pod.name,
+        }
diff --git a/robusta_krr/core/models/severity.py b/robusta_krr/core/models/severity.py
index d2c942e..ec1fad0 100644
--- a/robusta_krr/core/models/severity.py
+++ b/robusta_krr/core/models/severity.py
@@ -39,28 +39,28 @@ class Severity(str, enum.Enum):
 
         return calculate_severity(current, recommended, resource_type)
 
-    @staticmethod
-    def bind_calculator(resource_type: ResourceType) -> Callable[[SeverityCalculator], SeverityCalculator]:
-        """
-        Bind a severity calculator function to a resource type.
-        Use this decorator to bind a severity calculator function to a resource type.
-
-        Example:
-        >>> @bind_calculator(ResourceType.CPU)
-        >>> def cpu_severity_calculator(current: Optional[float], recommended: Optional[float], resource_type: ResourceType) -> Severity:
-        >>>     if current is None and recommended is None:
-        >>>         return Severity.GOOD
-        >>>     if current is None or recommended is None:
-        >>>         return Severity.WARNING
-        >>>
-        >>>     return Severity.CRITICAL if abs(current - recommended) >= 0.5 else Severity.GOOD
-        """
-
-        def decorator(func: SeverityCalculator) -> SeverityCalculator:
-            SEVERITY_CALCULATORS_REGISTRY[resource_type] = func
-            return func
-
-        return decorator
+
+def register_severity_calculator(resource_type: ResourceType) -> Callable[[SeverityCalculator], SeverityCalculator]:
+    """
+    Bind a severity calculator function to a resource type.
+    Use this decorator to bind a severity calculator function to a resource type.
+
+    Example:
+    >>> @bind_severity_calculator(ResourceType.CPU)
+    >>> def cpu_severity_calculator(current: Optional[float], recommended: Optional[float], resource_type: ResourceType) -> Severity:
+    >>>     if current is None and recommended is None:
+    >>>         return Severity.GOOD
+    >>>     if current is None or recommended is None:
+    >>>         return Severity.WARNING
+    >>>
+    >>>     return Severity.CRITICAL if abs(current - recommended) >= 0.5 else Severity.GOOD
+    """
+
+    def decorator(func: SeverityCalculator) -> SeverityCalculator:
+        SEVERITY_CALCULATORS_REGISTRY[resource_type] = func
+        return func
+
+    return decorator
 
 
 SeverityCalculator = Callable[[Optional[float], Optional[float], ResourceType], Severity]
@@ -86,7 +86,7 @@ def default_severity_calculator(
     return Severity.UNKNOWN
 
 
-@Severity.bind_calculator(ResourceType.CPU)
+@register_severity_calculator(ResourceType.CPU)
 def cpu_severity_calculator(
     current: Optional[float], recommended: Optional[float], resource_type: ResourceType
 ) -> Severity:
@@ -107,7 +107,7 @@ def cpu_severity_calculator(
         return Severity.GOOD
 
 
-@Severity.bind_calculator(ResourceType.Memory)
+@register_severity_calculator(ResourceType.Memory)
 def memory_severity_calculator(
     current: Optional[float], recommended: Optional[float], resource_type: ResourceType
 ) -> Severity:
diff --git a/robusta_krr/formatters/json.py b/robusta_krr/formatters/json.py
index 2c3a51e..dfded37 100644
--- a/robusta_krr/formatters/json.py
+++ b/robusta_krr/formatters/json.py
@@ -18,4 +18,4 @@ class JSONFormatter(BaseFormatter):
         :rtype: str
         """
 
-        return result.json(indent=2)
+        return result.json(indent=2, models_as_dict=False)