Repository: yasufumy/pipelib

Issue:
Fix bug in Dataset.save
```py
class Dataset:
    def save(self, filename):
        evaluated_dataset = list(self)
        with open(filename, 'wb') as f:
            pickle.dump(evaluated_dataset, f)
        return EvaluatedDataset(evaluated_dataset)

class EvaluatedDataset(PipelinedDataset):
    ...
```

Discussion:

Patch:
diff --git a/pipelib/core.py b/pipelib/core.py
index ec31d3e..2329b66 100644
--- a/pipelib/core.py
+++ b/pipelib/core.py
@@ -4,14 +4,13 @@ from pathlib import Path
 from itertools import chain, islice, tee
 from collections import deque
 
-import pipelib
 from pipelib import parallel
 from pipelib import iterators
 
 
 class Dataset:
     def __init__(self, dataset):
-        if isinstance(dataset, pipelib.Dataset):
+        if isinstance(dataset, Dataset):
             self._dataset = dataset._dataset
         else:
             self._dataset = dataset
@@ -67,14 +66,14 @@ class Dataset:
         return PipelinedDataset(self, f)
 
     def zip(self, *others):
-        assert all(isinstance(other, pipelib.Dataset) for other in others)
+        assert all(isinstance(other, Dataset) for other in others)
 
         def f(dataset):
             yield from zip(dataset, *others)
         return PipelinedDataset(self, f)
 
     def concat(self, *others):
-        assert all(isinstance(other, pipelib.Dataset) for other in others)
+        assert all(isinstance(other, Dataset) for other in others)
 
         def f(dataset):
             yield from chain(dataset, *others)
@@ -102,10 +101,10 @@ class Dataset:
         return next(iter(self))
 
     def save(self, filename):
-        evaluated_dataset = list(self)
+        cache = list(self)
         with open(filename, 'wb') as f:
-            pickle.dump(evaluated_dataset, f)
-        return self
+            pickle.dump(cache, f)
+        return CacheDataset(self, cache)
 
     @staticmethod
     def load(filename):
@@ -137,7 +136,7 @@ class _NestedFunc:
 class PipelinedDataset(Dataset):
 
     def __init__(self, dataset, func):
-        if not isinstance(dataset, pipelib.core.PipelinedDataset):
+        if not isinstance(dataset, PipelinedDataset):
             self._func = func
         else:
             self._func = _NestedFunc(dataset._func, func)
@@ -148,6 +147,16 @@ class PipelinedDataset(Dataset):
         yield from self._func(self._dataset)
 
 
+class CacheDataset(PipelinedDataset):
+
+    def __init__(self, dataset, cache):
+        super().__init__(dataset, lambda x: x)
+        self._cache = cache
+
+    def __iter__(self):
+        yield from self._cache
+
+
 class _Repeated:
     __slots__ = ['_generator', '_args', '_kwargs']
 
diff --git a/setup.py b/setup.py
index 75a9f8e..f4b0f3d 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@ except ImportError:
 
 setup(
     name='pipelib',
-    version='0.2.3',
+    version='0.2.4',
     description='minimal data pipeline library for Python',
     long_description=open('./README.md', encoding='utf-8').read(),
     long_description_content_type='text/markdown',


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
index 41731c9..f76895a 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -221,6 +221,14 @@ class DatasetTestCase(TestCase):
         expected = [x ** 2 for x in self.base if x % 2 == 0]
         self.assertListEqual(data.all(), expected)
         self.check_correct_pipelined_dataset(data, self.base)
+        self.assertIsInstance(data, pipelib.core.CacheDataset)
+
+        data = data.map(lambda x: x ** 2)
+        result = list(data._func(self.base))
+        expected = [x ** 2 for x in expected]
+        self.assertListEqual(data.all(), expected)
+        self.assertListEqual(result, expected)
+        self.check_correct_pipelined_dataset(data, self.base)
 
     @patch('pipelib.core.open')
     @patch('pipelib.core.pickle.load')