From 76d88df1187960f1056c0bad1d8621320753c207 Mon Sep 17 00:00:00 2001
From: Rambaud Pierrick <12rambau@users.noreply.github.com>
Date: Tue, 17 Dec 2024 07:22:22 +0000
Subject: [PATCH] feat: add rounding capabilities to data_regression

---
 src/pytest_regressions/common.py              | 22 +++++++++++++++++++
 src/pytest_regressions/data_regression.py     |  7 ++++++
 tests/test_data_regression.py                 |  9 ++++++++
 .../test_round_digits.yml                     |  4 ++++
 4 files changed, 42 insertions(+)
 create mode 100644 tests/test_data_regression/test_round_digits.yml

diff --git a/src/pytest_regressions/common.py b/src/pytest_regressions/common.py
index 30f118d..e3f13c7 100644
--- a/src/pytest_regressions/common.py
+++ b/src/pytest_regressions/common.py
@@ -4,6 +4,7 @@
 from typing import Callable
 from typing import List
 from typing import Optional
+from typing import Union
 
 import pytest
 
@@ -188,3 +189,24 @@ def make_location_message(banner: str, filename: Path, aux_files: List[str]) ->
             else:
                 dump_aux_fn(Path(obtained_filename))
                 raise
+
+
+def round_digits(data: Union[list, dict], prescision: int) -> Union[list, dict]:
+    """
+    Recusrsively Round the values of any float value in a collection to the given prescision.
+
+    :param data: The collection to round.
+    :param prescision: The number of decimal places to round to.
+    :return: The collection with all float values rounded to the given prescision.
+
+    """
+    # change the generator depending on the collection type
+    generator = enumerate(data) if isinstance(data, list) else data.items()
+    for k, v in generator:
+        if isinstance(v, (list, dict)):
+            data[k] = round_digits(v, prescision)
+        elif isinstance(v, float):
+            data[k] = round(v, prescision)
+        else:
+            data[k] = v
+    return data
diff --git a/src/pytest_regressions/data_regression.py b/src/pytest_regressions/data_regression.py
index 70a1d09..46874ee 100644
--- a/src/pytest_regressions/data_regression.py
+++ b/src/pytest_regressions/data_regression.py
@@ -11,6 +11,7 @@
 
 from .common import check_text_files
 from .common import perform_regression_check
+from .common import round_digits
 
 
 class DataRegressionFixture:
@@ -32,6 +33,7 @@ def check(
         data_dict: Dict[str, Any],
         basename: Optional[str] = None,
         fullpath: Optional["os.PathLike[str]"] = None,
+        prescision: Optional[int] = None,
     ) -> None:
         """
         Checks the given dict against a previously recorded version, or generate a new file.
@@ -46,10 +48,15 @@ def check(
             will ignore ``datadir`` fixture when reading *expected* files but will still use it to
             write *obtained* files. Useful if a reference file is located in the session data dir for example.
 
+        :param prescision: if given, round all floats in the dict to the given number of digits.
+
         ``basename`` and ``fullpath`` are exclusive.
         """
         __tracebackhide__ = True
 
+        if prescision is not None:
+            data_dict = round_digits(data_dict, prescision)
+
         def dump(filename: Path) -> None:
             """Dump dict contents to the given filename"""
             import yaml
diff --git a/tests/test_data_regression.py b/tests/test_data_regression.py
index c36e003..d79b5cf 100644
--- a/tests/test_data_regression.py
+++ b/tests/test_data_regression.py
@@ -38,6 +38,15 @@ def dump_scalar(dumper, scalar):
     data_regression.check(contents)
 
 
+def test_round_digits(data_regression):
+    """Example including float numbers and check rounding capabilities."""
+    contents = {
+        "content": {"value1": "toto", "value": 1.123456789},
+        "value": 1.23456789,
+    }
+    data_regression.check(contents, prescision=2)
+
+
 def test_usage_workflow(pytester, monkeypatch):
     monkeypatch.setattr(
         sys, "testing_get_data", lambda: {"contents": "Foo", "value": 10}, raising=False
diff --git a/tests/test_data_regression/test_round_digits.yml b/tests/test_data_regression/test_round_digits.yml
new file mode 100644
index 0000000..ea5f70a
--- /dev/null
+++ b/tests/test_data_regression/test_round_digits.yml
@@ -0,0 +1,4 @@
+content:
+  value: 1.12
+  value1: toto
+value: 1.23