From d30a65e8cd8f22816dae31fa5b9ac3650764a48b Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Fri, 20 Dec 2019 11:08:57 +0100
Subject: [PATCH] fix return type

---
 pts/dataset/common.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/pts/dataset/common.py b/pts/dataset/common.py
index c4440f3..45ab4e5 100644
--- a/pts/dataset/common.py
+++ b/pts/dataset/common.py
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Iterable, NamedTuple, Sized, List, Optional
+from typing import Any, Dict, Iterable, NamedTuple, Sized, List, Optional, Iterator
 
 from pydantic import BaseModel
 
@@ -38,20 +38,23 @@ class FieldName:
 
 class Dataset(Sized, Iterable[DataEntry], ABC):
     @abstractmethod
-    def __iter__(self) -> Iterable[DataEntry]:
+    def __iter__(self) -> Iterator[DataEntry]:
         pass
 
     @abstractmethod
     def __len__(self):
         pass
 
+
 class CategoricalFeatureInfo(BaseModel):
     name: str
     cardinality: str
 
+
 class BasicFeatureInfo(BaseModel):
     name: str
 
+
 class MetaData(BaseModel):
     freq: str = None
     target: Optional[BasicFeatureInfo] = None
@@ -73,4 +76,3 @@ class TrainDatasets(NamedTuple):
     metadata: MetaData
     train: Dataset
     test: Optional[Dataset] = None
-