Add return type to dataclassreader

reimplement dfurtado/dataclass-csv#57

Author: kraktus

dataclass_csv/dataclass_reader.py

@@ -2,7 +2,7 @@
 import csv
 
 from datetime import date, datetime
-from typing import Union, Type, Optional, Sequence, Dict, Any, List
+from typing import Union, Type, Optional, Sequence, Dict, Any, List, Generic, TypeVar
 
 import typing
 
@@ -12,6 +12,8 @@
 from collections import Counter
 
 
+T = TypeVar("T")
+
 def _verify_duplicate_header_items(header):
     if header is not None and len(header) == 0:
         return
@@ -60,11 +62,11 @@ def get_args(t):
     return tuple()
 
 
-class DataclassReader:
+class DataclassReader(Generic[T]):
     def __init__(
         self,
         f: Any,
-        cls: Type[object],
+        cls: Type[T],
         fieldnames: Optional[Sequence[str]] = None,
         restkey: Optional[str] = None,
         restval: Optional[Any] = None,
@@ -198,7 +200,7 @@ def _parse_date_value(self, field, date_value, field_type):
         else:
             return datetime_obj
 
-    def _process_row(self, row):
+    def _process_row(self, row) -> T:
         values = dict()
 
         for field in dataclasses.fields(self._cls):
@@ -257,7 +259,7 @@ def _process_row(self, row):
                 values[field.name] = transformed_value
         return self._cls(**values)
 
-    def __next__(self):
+    def __next__(self) -> T:
         row = next(self._reader)
         return self._process_row(row)
 

dataclass_csv/dataclass_writer.py

@@ -11,7 +11,7 @@ def __init__(
         data: List[Any],
         cls: Type[object],
         dialect: str = "excel",
-        **fmtparams: Dict[str, Any],
+        **fmtparams: Any,
     ):
         if not f:
             raise ValueError("The f argument is required")