add __str__ method

Serapieum-of-alex · Aug 17, 2024 · b70deff · b70deff
1 parent 1ab4717
commit b70deff
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/statista/distributions.py b/statista/distributions.py
@@ -4,6 +4,7 @@
 from typing import Any, List, Tuple, Union, Dict, Callable
 from abc import ABC, abstractmethod
 import numpy as np
+from statistics import mode
 import scipy.optimize as so
 from matplotlib.figure import Figure
 from matplotlib.axes import Axes
@@ -144,6 +145,27 @@ def __init__(
         else:
             raise TypeError("The `parameters` argument should be dictionary")
 
+    def __str__(self) -> str:
+        message = ""
+        if self.data is not None:
+            message += f"""
+                    Dataset of {len(self.data)} value
+                    min: {np.min(self.data)}
+                    max: {np.max(self.data)}
+                    mean: {np.mean(self.data)}
+                    median: {np.median(self.data)}
+                    mode: {mode(self.data)}
+                    std: {np.std(self.data)}
+                    Distribution : {self.__class__.__name__}
+                    parameters: {self.parameters}
+                    """
+        if self.parameters is not None:
+            message += f"""
+                Distribution : {self.__class__.__name__}
+                parameters: {self.parameters}
+                """
+        return message
+
     @property
     def parameters(self) -> Dict[str, float]:
         """Distribution parameters"""

diff --git a/tests/test_distributions.py b/tests/test_distributions.py
@@ -39,6 +39,24 @@ def test_plotting_position_rp(
         assert isinstance(rp, np.ndarray)
 
 
+class TestAbstractDistribution:
+    def test_abstract_distribution(self, time_series1: list, gev_dist_parameters):
+        text_1 = "\n                    Dataset of 27 value\n                    min: 15.790480003140171\n                    max: 19.39645340792385\n                    mean: 16.929171461473548\n                    median: 16.626465201654593\n                    mode: 15.999737471905252\n                    std: 1.0211514099144634\n                    Distribution : Gumbel\n                    parameters: None\n                    "
+        parameters = gev_dist_parameters["lmoments"]
+        dist = Gumbel(time_series1)
+        assert str(dist) == text_1
+
+        text_2 = (
+            "\n                Distribution : Gumbel\n                parameters: {'loc': 16.392889171307772, "
+            "'scale': 0.7005442761744839, 'shape': -0.1614793298009645}\n                "
+        )
+        dist = Gumbel(parameters=parameters)
+        assert str(dist) == text_2
+        dist = Gumbel(data=time_series1, parameters=parameters)
+        text_3 = "\n                    Dataset of 27 value\n                    min: 15.790480003140171\n                    max: 19.39645340792385\n                    mean: 16.929171461473548\n                    median: 16.626465201654593\n                    mode: 15.999737471905252\n                    std: 1.0211514099144634\n                    Distribution : Gumbel\n                    parameters: {'loc': 16.392889171307772, 'scale': 0.7005442761744839, 'shape': -0.1614793298009645}\n                    \n                Distribution : Gumbel\n                parameters: {'loc': 16.392889171307772, 'scale': 0.7005442761744839, 'shape': -0.1614793298009645}\n                "
+        assert str(dist) == text_3
+
+
 class TestGumbel:
     def test_create_instance(
         self,