Merge pull request #254 from tofunlp/feature/return-slice-iterabledat…

…aset

Subset (slice) of iterabledataset should be iterabledataset, not list.

lineflow/core.py

@@ -206,6 +206,12 @@ def get_example(self, i: int) -> Any:
     def __len__(self) -> int:
         return super(IterableDataset, self).__len__()
 
+    def __getitem__(self, index: Union[int, slice]) -> Union[Any, List[Any]]:
+        out = super(IterableDataset, self).__getitem__(index)
+        if isinstance(out, list):
+            out = Dataset(out)
+        return out
+
 
 class ConcatDataset(Dataset):
     def __init__(self, *datasets: List[DatasetMixin]) -> None:

tests/test_core.py

@@ -158,6 +158,14 @@ def test_dunder_len(self):
         self.assertEqual(self.data._length, len(self.base))
         self.assertTrue(self.data._computed)
 
+    def test_dunder_slice(self):
+        # Subset of IterableDataset should be Dataset.
+        subset = self.data[:10]
+        self.assertIsInstance(subset, Dataset)
+
+        single_sample = self.data[0]
+        self.assertIsInstance(single_sample, int)
+
 
 class DatasetTestCase(TestCase):