Make database implementations iterable

JohnVinyard · May 13, 2016 · 1f85a9f · 1f85a9f
1 parent eb8d54d
commit 1f85a9f
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 1 deletion.
diff --git a/featureflow/__init__.py b/featureflow/__init__.py
@@ -1,4 +1,4 @@
-__version__ = '0.7.9'
+__version__ = '0.8.9'
 
 from model import BaseModel
 

diff --git a/featureflow/data.py b/featureflow/data.py
@@ -102,6 +102,9 @@ def size(self, key):
     def iter_ids(self):
         raise NotImplementedError()
 
+    def __iter__(self):
+        return self.iter_ids()
+
     def __contains__(self, key):
         raise NotImplementedError()
 

diff --git a/featureflow/test_integration.py b/featureflow/test_integration.py
@@ -306,6 +306,19 @@ class MultipleRoots(BaseModel):
 
 
 class BaseTest(object):
+
+    def test_can_iterate_over_database(self):
+        class D(BaseModel, self.Settings):
+            stream = Feature(TextStream, store=True)
+            words = Feature(Tokenizer, needs=stream, store=False)
+            count = JSONFeature(WordCount, needs=words, store=True)
+
+        D.process(stream='mary')
+        D.process(stream='humpty')
+        D.process(stream='cased')
+
+        self.assertEqual(3, len(list(self.Settings.database)))
+
     def test_can_use_iterator_node(self):
         iterable = chunked(StringIO(data_source['mary']), chunksize=3)