Repository: tofunlp/lineflow

Issue:
Should slice of IterableDataset return IterableDataset not List?
**Is your feature request related to a problem? Please describe.**

```py
train = lfds.SciTLDR(split="train")  # IterableDataset
train_mini = train[:10]  # Now this is just a python list (List[Any])
```

If I make a subset of a dataset, it loses all the features such as `.map`.

**Describe the solution you'd like**
Return `IterableDataset` in stead of List.

Discussion:

Patch:
diff --git a/lineflow/core.py b/lineflow/core.py
index b9d31c7..8bcfb0a 100644
--- a/lineflow/core.py
+++ b/lineflow/core.py
@@ -206,6 +206,12 @@ class IterableDataset(Dataset):
     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:


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
index 6171a5a..9e40110 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -158,6 +158,14 @@ class IterableDatasetTestCase(TestCase):
         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):