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},