Repository: sotetsuk/memozo

Issue:
implement generator
See #2 

Discussion:

Patch:
diff --git a/memozo/memozo.py b/memozo/memozo.py
index 4832552..a67bbe5 100644
--- a/memozo/memozo.py
+++ b/memozo/memozo.py
@@ -40,3 +40,39 @@ class Memozo(object):
             return _wrapper
 
         return wrapper
+
+    def generator(self, name=None, ext='file'):
+
+        def wrapper(func):
+            _name = func.__name__ if name is None else name
+
+            @functools.wraps(func)
+            def _wrapper(*args, **kwargs):
+                # get cached data path
+                args = utils.get_bound_args(func, *args, **kwargs)
+                args_str = utils.get_args_str(args)
+                sha1 = utils.get_hash(_name, func.__name__, args_str)
+                file_path = os.path.join(self.base_path, "{}_{}.{}".format(_name, sha1, ext))
+
+                # if cached data exists, return generator using cached data
+                if utils.log_exisits(self.base_path, _name, func.__name__, args_str) and os.path.exists(file_path):
+                    def gen_cached_data():
+                        with codecs.open(file_path, 'r', utils.ENCODING) as f:
+                            for line in f:
+                                yield line
+                    return gen_cached_data
+
+                gen = func(*args, **kwargs)
+
+                # if no cached data exists, generator not only yield value but save value at each iteration
+                def generator_with_cache(gen, file_path):
+                    with codecs.open(file_path, 'w', utils.ENCODING) as f:
+                        for e in gen:
+                            f.write(e)
+                            yield e
+
+                return generator_with_cache(gen, file_path)
+
+            return _wrapper
+
+        return wrapper


Test patch:
diff --git a/tests/test_memozo.py b/tests/test_memozo.py
index b47b5d2..41b3642 100644
--- a/tests/test_memozo.py
+++ b/tests/test_memozo.py
@@ -1,5 +1,6 @@
 import os
 import unittest
+import codecs
 
 from memozo import Memozo, utils
 
@@ -70,4 +71,41 @@ class TestMemozoCall(unittest.TestCase):
         create_dummy_data(3)
         self.assertTrue(os.path.exists(file_path))
 
-        os.remove(file_path)
\ No newline at end of file
+        os.remove(file_path)
+
+
+class TestMemozoGenerator(unittest.TestCase):
+
+    def test_no_cache_output(self):
+        base_path = './tests/resources'
+        m = Memozo(base_path)
+
+        @m.generator('gen_test')
+        def gen_test_func():
+            for i in range(10):
+                if i % 3 == 0:
+                    yield "{}\n".format(i)
+
+        gen = gen_test_func()
+        for i in gen:
+            self.assertTrue(int(i.strip('\n')) % 3 == 0)
+
+    def test_data_cached_collectly(self):
+        base_path = './tests/resources'
+        m = Memozo(base_path)
+        sha1 = utils.get_hash('gen_test', 'gen_test_func', '')
+        file_path = os.path.join(base_path, "{}_{}.{}".format('gen_test', sha1, 'file'))
+
+        @m.generator('gen_test')
+        def gen_test_func():
+            for i in range(10):
+                if i % 3 == 0:
+                    yield "{}\n".format(i)
+
+        gen1 = gen_test_func()
+        for _ in gen1:
+            continue
+
+        with codecs.open(file_path, 'r', 'utf-8') as f:
+            for line in f:
+                self.assertTrue(int(line.strip('\n')) % 3 == 0)