From 582c52fa8cdf82df256c3ecf1e7a3a351a9e5f32 Mon Sep 17 00:00:00 2001
From: Evgeny Ivanov <ivanov.evgeny.n@gmail.com>
Date: Sun, 15 Sep 2024 11:49:05 +0500
Subject: [PATCH] Support pretty formatting for ExperimentResults

---
 src/tea_tasting/experiment.py | 24 +++++++++++++++++++++++-
 tests/test_experiment.py      | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/src/tea_tasting/experiment.py b/src/tea_tasting/experiment.py
index 0919ffb..c81a21f 100644
--- a/src/tea_tasting/experiment.py
+++ b/src/tea_tasting/experiment.py
@@ -306,7 +306,29 @@ def to_html(
         """
         return tea_tasting.utils.PrettyDictsMixin.to_html(self, keys, formatter)
 
-ExperimentResults = dict[tuple[Any, Any], ExperimentResult]
+
+class ExperimentResults(
+    UserDict[tuple[Any, Any], ExperimentResult],
+    tea_tasting.utils.PrettyDictsMixin,
+):
+    """Experiment results for multiple pairs of variants."""
+    default_keys = (
+        "variants",
+        "metric",
+        "control",
+        "treatment",
+        "rel_effect_size",
+        "rel_effect_size_ci",
+        "pvalue",
+    )
+
+    def to_dicts(self) -> tuple[dict[str, Any], ...]:
+        """Convert the result to a sequence of dictionaries."""
+        return tuple(
+            {"variants": str(variants)} | metric_result
+            for variants, experiment_result in self.items()
+            for metric_result in experiment_result.to_dicts()
+        )
 
 
 class ExperimentPowerResult(
diff --git a/tests/test_experiment.py b/tests/test_experiment.py
index c6f5ab3..681fb19 100644
--- a/tests/test_experiment.py
+++ b/tests/test_experiment.py
@@ -316,6 +316,41 @@ def test_experiment_result_to_html(result2: tea_tasting.experiment.ExperimentRes
     )).to_html(index=False)
 
 
+def test_experiment_results_to_dicts(
+    results2: tea_tasting.experiment.ExperimentResults,
+):
+    assert results2.to_dicts() == (
+        {
+            "variants": "(0, 1)",
+            "metric": "metric_tuple",
+            "control": 10,
+            "treatment": 11,
+            "effect_size": 1,
+        },
+        {
+            "variants": "(0, 1)",
+            "metric": "metric_dict",
+            "control": 20,
+            "treatment": 22,
+            "effect_size": 2,
+        },
+        {
+            "variants": "(0, 2)",
+            "metric": "metric_tuple",
+            "control": 10,
+            "treatment": 11,
+            "effect_size": 1,
+        },
+        {
+            "variants": "(0, 2)",
+            "metric": "metric_dict",
+            "control": 30,
+            "treatment": 33,
+            "effect_size": 3,
+        },
+    )
+
+
 def test_experiment_power_result_to_dicts():
     raw_results = (
         {"power": 0.8, "effect_size": 1, "rel_effect_size": 0.05, "n_obs": 20_000},