From ca1cee2f5283d4cdeb881a07254fbd32e1530e5e Mon Sep 17 00:00:00 2001
From: Matt Stone <matt@fulcrumgenomics.com>
Date: Wed, 16 Oct 2024 14:54:37 -0400
Subject: [PATCH] feat: support frozen dataclasses

---
 dataclass_io/_lib/dataclass_extensions.py | 14 +++++++-
 tests/test_reader.py                      | 43 +++++++++++------------
 2 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/dataclass_io/_lib/dataclass_extensions.py b/dataclass_io/_lib/dataclass_extensions.py
index d9b64c3..f1f84a8 100644
--- a/dataclass_io/_lib/dataclass_extensions.py
+++ b/dataclass_io/_lib/dataclass_extensions.py
@@ -56,11 +56,23 @@ def row_to_dataclass(
         # version of the dataclass with validation. We instantiate from this version to take
         # advantage of pydantic's validation, but then unpack the validated data in order to return
         # an instance of the user-specified dataclass.
-        pydantic_cls = pydantic_dataclass(dataclass_type)
+
+        params = dataclass_type.__dataclass_params__  # type:ignore[attr-defined]
+
+        pydantic_cls = pydantic_dataclass(
+            _cls=dataclass_type,
+            repr=params.repr,
+            eq=params.eq,
+            order=params.order,
+            unsafe_hash=params.unsafe_hash,
+            frozen=params.frozen,
+        )
+
         validated_data = pydantic_cls(**row)
         unpacked_data = {
             field.name: getattr(validated_data, field.name) for field in fields(dataclass_type)
         }
+
         data = dataclass_type(**unpacked_data)
 
     return data
diff --git a/tests/test_reader.py b/tests/test_reader.py
index 6a36d18..8a58e17 100644
--- a/tests/test_reader.py
+++ b/tests/test_reader.py
@@ -1,37 +1,36 @@
 from dataclasses import dataclass
 from pathlib import Path
+from typing import cast
 
-from dataclass_io.reader import DataclassReader
-
+import pytest
 
-@dataclass(kw_only=True, eq=True)
-class FakeDataclass:
-    foo: str
-    bar: int
+from dataclass_io.reader import DataclassReader
 
 
-def test_reader(tmp_path: Path) -> None:
+@pytest.mark.parametrize("kw_only", [True, False])
+@pytest.mark.parametrize("eq", [True, False])
+@pytest.mark.parametrize("frozen", [True, False])
+def test_reader(kw_only: bool, eq: bool, frozen: bool, tmp_path: Path) -> None:
     fpath = tmp_path / "test.txt"
 
+    @dataclass(frozen=frozen, eq=eq, kw_only=kw_only)  # type: ignore[literal-required]
+    class FakeDataclass:
+        foo: str
+        bar: int
+
     with fpath.open("w") as f:
         f.write("foo\tbar\n")
         f.write("abc\t1\n")
 
+    rows: list[FakeDataclass]
     with DataclassReader.open(filename=fpath, dataclass_type=FakeDataclass) as reader:
-        rows = [row for row in reader]
-
-    assert rows[0] == FakeDataclass(foo="abc", bar=1)
-
-
-def test_reader_from_str(tmp_path: Path) -> None:
-    """Test that we can create a reader when `filename` is a `str`."""
-    fpath = tmp_path / "test.txt"
-
-    with fpath.open("w") as f:
-        f.write("foo\tbar\n")
-        f.write("abc\t1\n")
+        # TODO make `DataclassReader` generic
+        rows = cast(list[FakeDataclass], [row for row in reader])
 
-    with DataclassReader.open(filename=str(fpath), dataclass_type=FakeDataclass) as reader:
-        rows = [row for row in reader]
+    assert len(rows) == 1
 
-    assert rows[0] == FakeDataclass(foo="abc", bar=1)
+    if eq:
+        assert rows[0] == FakeDataclass(foo="abc", bar=1)
+    else:
+        assert rows[0].foo == "abc"
+        assert rows[0].bar == 1