Support non-list iterable as input to DataclassWriter

reimplement and close dfurtado/dataclass-csv#61

Author: kraktus

dataclass_csv/dataclass_writer.py

@@ -1,24 +1,25 @@
 import csv
 import dataclasses
-from typing import Type, Dict, Any, List
+from typing import Type, Dict, Any, List, Iterable, Generic, TypeVar
 from .header_mapper import HeaderMapper
 
 
-class DataclassWriter:
+
+T = TypeVar("T")
+
+
+class DataclassWriter(Generic[T]):
     def __init__(
         self,
         f: Any,
-        data: List[Any],
-        klass: Type[object],
+        data: Iterable[T],
+        klass: Type[T],
         dialect: str = "excel",
         **fmtparams: Any,
     ):
         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")
-
         if not dataclasses.is_dataclass(klass):
             raise ValueError("Invalid 'klass' argument. It must be a dataclass")
 
@@ -48,7 +49,6 @@ def write(self, skip_header: bool = False):
                 self._fieldnames = self._apply_mapping()
 
             self._writer.writerow(self._fieldnames)
-
         for item in self._data:
             if not isinstance(item, self._cls):
                 raise TypeError(

tests/test_dataclass_writer.py

@@ -61,12 +61,16 @@ def test_invalid_file_value(tmpdir_factory):
     with pytest.raises(ValueError):
         DataclassWriter(None, users, User)
 
-
-def test_with_data_not_a_list(tmpdir_factory):
+def test_with_iterable(tmpdir_factory):
     tempfile = tmpdir_factory.mktemp("data").join("user_001.csv")
-
-    users = User(name="test", age=40)
+    users_dict = {"test": User(name="test", age=40)}
 
     with tempfile.open("w") as f:
-        with pytest.raises(ValueError):
-            DataclassWriter(f, users, User)
+        DataclassWriter(f, users_dict.values(), User).write()
+
+    with tempfile.open() as f:
+        reader = DataclassReader(f, User)
+        saved_users = list(reader)
+
+        assert len(saved_users) > 0
+        assert saved_users[0].name == users_dict["test"].name