Add a context manager for modifying the JSON encoder

cfspopcon/file_io.py

@@ -2,7 +2,7 @@
 
 import json
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, Self
 
 import numpy as np
 import xarray as xr
@@ -91,6 +91,38 @@ def read_dataset_from_netcdf(filepath: Path) -> xr.Dataset:
     return dataset
 
 
+# These following lines are needed to modify the representation of floats
+# in the JSON file. This is needed, because otherwise small errors in the
+# floats lead to a large, meaningless diff of the reference JSON files.
+
+
+class RoundingFloat(float):
+    """A modified version of the 'float' built-in, with a modified __repr__.
+
+    This is needed because `iterencode` directly uses `float.__repr__` in `floatstr`.
+
+    From: https://stackoverflow.com/questions/54370322/how-to-limit-the-number-of-float-digits-jsonencoder-produces
+    """
+
+    __repr__ = staticmethod(lambda x: f"{x:#.10g}")  # type:ignore[assignment,unused-ignore]
+
+
+class ModifyJSONFloatRepr:
+    """A ContextManager to locally modify the representation of floats."""
+
+    def __enter__(self) -> Self:
+        """Change the float representation to fixed precision on entry."""
+        self.c_make_encoder = json.encoder.c_make_encoder  # type:ignore[attr-defined]
+        json.encoder.c_make_encoder = None  # type:ignore[attr-defined]
+        json.encoder.float = RoundingFloat  # type:ignore[attr-defined]
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        """Change the float representation back to the default."""
+        json.encoder.c_make_encoder = self.c_make_encoder  # type:ignore[attr-defined]
+        json.encoder.float = float  # type:ignore[attr-defined]
+
+
 def write_point_to_file(dataset: xr.Dataset, point_key: str, point_params: dict, output_dir: Path) -> None:
     """Write the analysis values at the named points to a json file."""
     mask = build_mask_from_dict(dataset, point_params)
@@ -121,16 +153,6 @@ def write_point_to_file(dataset: xr.Dataset, point_key: str, point_params: dict,
 
     output_dir.mkdir(parents=True, exist_ok=True)
 
-    class RoundingFloat(float):
-        """A formatter to control how floats are written to JSON.
-
-        From: https://stackoverflow.com/questions/54370322/how-to-limit-the-number-of-float-digits-jsonencoder-produces
-        """
-
-        __repr__ = staticmethod(lambda x: f"{x:#.10g}")  # type:ignore[assignment,unused-ignore]
-
-    json.encoder.c_make_encoder = None  # type:ignore[attr-defined]
-    json.encoder.float = RoundingFloat  # type:ignore[attr-defined]
-
-    with open(output_dir / f"{point_key}.json", "w") as file:
-        json.dump(point.to_dict(), file, indent=4, sort_keys=True)
+    with ModifyJSONFloatRepr():
+        with open(output_dir / f"{point_key}.json", "w") as file:
+            json.dump(point.to_dict(), file, indent=4, sort_keys=True)