diff --git a/fsql/deser.py b/fsql/deser.py index a06dfdc..327622d 100644 --- a/fsql/deser.py +++ b/fsql/deser.py @@ -53,7 +53,7 @@ from enum import Enum, auto, unique from functools import partial, reduce from itertools import chain -from typing import Any, Generic, Iterable, NamedTuple, Tuple, TypeVar, Union +from typing import Any, Generic, Iterable, Tuple, TypeVar, Union import pandas as pd from fsspec.core import OpenFile @@ -98,12 +98,14 @@ def from_url(cls, url: str): DataObject = TypeVar("DataObject") -class DataObjectRich(NamedTuple, Generic[DataObject]): +@dataclass(frozen=True) +class DataObjectRich(Generic[DataObject]): data: DataObject failures: Iterable[PartitionReadFailure] -class PartitionReadFailure(NamedTuple): +@dataclass(frozen=True) +class PartitionReadFailure: partition: Partition reason: str # or Any? diff --git a/tests/test_dict_reader.py b/tests/test_dict_reader.py index 00012e7..da4679e 100644 --- a/tests/test_dict_reader.py +++ b/tests/test_dict_reader.py @@ -39,5 +39,5 @@ def test_lazy_errors(tmp_path): lazy_reader = EnumeratedDictReader(lazy_errors=True) result = read_partitioned_table(f"file://{case1_path}/", Q_TRUE, data_reader=lazy_reader) - assert result[0] == {0: json.loads(data1)} - assert [error_line] == [e.reason for e in result[1]] + assert result.data == {0: json.loads(data1)} + assert [error_line] == [e.reason for e in result.failures] diff --git a/tests/test_pandasreader.py b/tests/test_pandasreader.py index e6e3d1f..ec798eb 100644 --- a/tests/test_pandasreader.py +++ b/tests/test_pandasreader.py @@ -51,6 +51,6 @@ def test_lazy_errors(tmp_path): result = read_partitioned_table(f"file://{case1_path}/", Q_TRUE, data_reader=reader_eager) reader_lazy = PandasReader(columns=["c3"], lazy_errors=True) result = read_partitioned_table(f"file://{case1_path}/", Q_TRUE, data_reader=reader_lazy) - assert_frame_equal(df2[["c3"]], result[0]) - reasons = [e.reason.split("\n")[0] for e in result[1]] + assert_frame_equal(df2[["c3"]], result.data) + reasons = [e.reason.split("\n")[0] for e in result.failures] assert reasons == [error_line]