From 2a2a09d18204cb44250eff7b58d18ce457a94f5a Mon Sep 17 00:00:00 2001 From: Matt Stone Date: Wed, 4 Mar 2026 07:10:01 -0800 Subject: [PATCH 1/3] fix: Use JSON mode for model serialization in MetricWriter Use `model_dump(mode="json")` in `MetricWriter.write()` so that pydantic serializers run fully. Handle both `str` and `StrEnum` keys in `_pivot_counter_values` to support the already-converted keys that JSON mode produces. Co-Authored-By: Rahul Kaushal Co-Authored-By: Claude Opus 4.6 --- fgmetric/collections/_counter_pivot_table.py | 3 +- fgmetric/metric_writer.py | 2 +- tests/test_collections.py | 18 ++++++++++++ tests/test_metric_writer.py | 29 ++++++++++++++++++++ 4 files changed, 50 insertions(+), 2 deletions(-) diff --git a/fgmetric/collections/_counter_pivot_table.py b/fgmetric/collections/_counter_pivot_table.py index c839489..7855521 100644 --- a/fgmetric/collections/_counter_pivot_table.py +++ b/fgmetric/collections/_counter_pivot_table.py @@ -181,6 +181,7 @@ def _pivot_counter_values( # Replace the counter field with keys for each of its enum's members counts = data.pop(self._counter_fieldname) for key, count in counts.items(): - data[key.value] = count + column_name = key if isinstance(key, str) else key.value + data[column_name] = count return data diff --git a/fgmetric/metric_writer.py b/fgmetric/metric_writer.py index b84888d..a311fa3 100644 --- a/fgmetric/metric_writer.py +++ b/fgmetric/metric_writer.py @@ -106,7 +106,7 @@ def write(self, metric: T) -> None: TypeError: If the provided `metric` is not an instance of the Metric class used to parametrize the writer. """ - self._writer.writerow(metric.model_dump()) + self._writer.writerow(metric.model_dump(mode="json")) def writeall(self, metrics: Iterable[T]) -> None: """ diff --git a/tests/test_collections.py b/tests/test_collections.py index 0380522..b5bd379 100644 --- a/tests/test_collections.py +++ b/tests/test_collections.py @@ -204,6 +204,24 @@ class FakeMetric(Metric): next(f) +def test_counter_pivot_table_model_dump_json_mode() -> None: + """Test that model_dump(mode='json') works with Counter pivot tables.""" + + @unique + class FakeEnum(StrEnum): + FOO = "foo" + BAR = "bar" + + class FakeMetric(Metric): + name: str + counts: Counter[FakeEnum] + + metric = FakeMetric(name="test", counts=Counter({FakeEnum.FOO: 1, FakeEnum.BAR: 2})) + result = metric.model_dump(mode="json") + + assert result == {"name": "test", "foo": 1, "bar": 2} + + def test_counter_pivot_table_missing_enum_members_default_to_zero(tmp_path: Path) -> None: """Test that missing enum members in input default to 0.""" diff --git a/tests/test_metric_writer.py b/tests/test_metric_writer.py index 7a5f05c..feb71f2 100644 --- a/tests/test_metric_writer.py +++ b/tests/test_metric_writer.py @@ -1,3 +1,6 @@ +from collections import Counter +from enum import StrEnum +from enum import unique from pathlib import Path from typing import assert_type from unittest import mock @@ -47,3 +50,29 @@ def test_init_closes_file_on_failure(tmp_path: Path) -> None: MetricWriter(FakeMetric, fpath) assert real_fout.closed + + +def test_writer_with_counter_metric(tmp_path: Path) -> None: + """Test we can write a Counter metric through MetricWriter.""" + + @unique + class FakeEnum(StrEnum): + FOO = "foo" + BAR = "bar" + + class CounterMetric(Metric): + name: str + counts: Counter[FakeEnum] + + fpath = tmp_path / "test.txt" + + with MetricWriter(CounterMetric, fpath) as writer: + writer.write( + CounterMetric(name="test", counts=Counter({FakeEnum.FOO: 3, FakeEnum.BAR: 4})) + ) + + with fpath.open("r") as f: + assert next(f) == "name\tfoo\tbar\n" + assert next(f) == "test\t3\t4\n" + with pytest.raises(StopIteration): + next(f) From abd9f15629b2e3cbfb59942c67fb717f9aa6c07a Mon Sep 17 00:00:00 2001 From: Matt Stone Date: Sat, 7 Mar 2026 11:09:57 -0800 Subject: [PATCH 2/3] style: Fix ruff formatting in test_metric_writer Co-Authored-By: Rahul Kaushal Co-Authored-By: Claude Opus 4.6 --- tests/test_metric_writer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_metric_writer.py b/tests/test_metric_writer.py index feb71f2..84b1f74 100644 --- a/tests/test_metric_writer.py +++ b/tests/test_metric_writer.py @@ -67,9 +67,7 @@ class CounterMetric(Metric): fpath = tmp_path / "test.txt" with MetricWriter(CounterMetric, fpath) as writer: - writer.write( - CounterMetric(name="test", counts=Counter({FakeEnum.FOO: 3, FakeEnum.BAR: 4})) - ) + writer.write(CounterMetric(name="test", counts=Counter({FakeEnum.FOO: 3, FakeEnum.BAR: 4}))) with fpath.open("r") as f: assert next(f) == "name\tfoo\tbar\n" From 2e3720d53329310b6bb4f71b6278f78110ff3246 Mon Sep 17 00:00:00 2001 From: Matt Stone Date: Sat, 7 Mar 2026 11:21:01 -0800 Subject: [PATCH 3/3] doc: Note JSON serialization mode in MetricWriter.write() Co-Authored-By: Claude Opus 4.6 --- fgmetric/metric_writer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/fgmetric/metric_writer.py b/fgmetric/metric_writer.py index a311fa3..21de462 100644 --- a/fgmetric/metric_writer.py +++ b/fgmetric/metric_writer.py @@ -97,7 +97,9 @@ def write(self, metric: T) -> None: """ Write a single Metric instance to file. - The Metric is converted to a dictionary and then written using the underlying `DictWriter`. + The Metric is serialized using ``model_dump(mode="json")`` and then written using the + underlying `DictWriter`. JSON mode ensures that all field values (e.g., enums) are + converted to JSON-compatible types before writing. Args: metric: An instance of the specified Metric.