diff --git a/dataclass_csv/dataclass_writer.py b/dataclass_csv/dataclass_writer.py index fdcbb52..bb82d59 100644 --- a/dataclass_csv/dataclass_writer.py +++ b/dataclass_csv/dataclass_writer.py @@ -1,6 +1,6 @@ import csv import dataclasses -from typing import Type, Dict, Any, List +from typing import Type, Dict, Any, Iterable from .header_mapper import HeaderMapper @@ -8,7 +8,7 @@ class DataclassWriter: def __init__( self, f: Any, - data: List[Any], + data: Iterable, cls: Type[object], dialect: str = "excel", **fmtparams: Dict[str, Any], @@ -16,8 +16,10 @@ def __init__( if not f: raise ValueError("The f argument is required") - if not isinstance(data, list): - raise ValueError("Invalid 'data' argument. It must be a list") + try: + iter(data) + except TypeError: + raise ValueError("Invalid 'data' argument. It must be an iterable") if not dataclasses.is_dataclass(cls): raise ValueError("Invalid 'cls' argument. It must be a dataclass") @@ -54,7 +56,7 @@ def write(self, skip_header: bool = False): raise TypeError( ( f"The item [{item}] is not an instance of " - f"{self._cls.__name__}. All items on the list must be " + f"{self._cls.__name__}. All items in the iterable must be " "instances of the same type" ) ) diff --git a/tests/test_dataclass_writer.py b/tests/test_dataclass_writer.py index 4b54517..5972c26 100644 --- a/tests/test_dataclass_writer.py +++ b/tests/test_dataclass_writer.py @@ -62,7 +62,7 @@ def test_invalid_file_value(tmpdir_factory): DataclassWriter(None, users, User) -def test_with_data_not_a_list(tmpdir_factory): +def test_with_data_not_an_iterable(tmpdir_factory): tempfile = tmpdir_factory.mktemp("data").join("user_001.csv") users = User(name="test", age=40)