Fix incorrect initialization ShardedCorpus with a generator. Fix #1511 (

#1512)

Fix incorrect initialization ShardedCorpus with a generator. Fix #1511.

gensim/corpora/sharded_corpus.py

@@ -251,7 +251,8 @@ def __init__(self, output_prefix, corpus, dim=None,
     def init_shards(self, output_prefix, corpus, shardsize=4096, dtype=_default_dtype):
         """Initialize shards from the corpus."""
 
-        if not gensim.utils.is_corpus(corpus):
+        is_corpus, corpus = gensim.utils.is_corpus(corpus)
+        if not is_corpus:
             raise ValueError(
                 "Cannot initialize shards without a corpus to read from! (Got corpus type: {0})".format(type(corpus))
             )

gensim/test/test_sharded_corpus.py

@@ -256,6 +256,18 @@ def test_resize(self):
             fname = dataset._shard_name(n)
             self.assertTrue(os.path.isfile(fname))
 
+    def test_init_with_generator(self):
+
+        def data_generator():
+            yield [(0, 1)]
+            yield [(1, 1)]
+
+        gen_tmp_fname = self.tmp_fname + '.generator'
+        corpus = ShardedCorpus(gen_tmp_fname, data_generator(), dim=2)
+
+        self.assertEqual(2, len(corpus))
+        self.assertEqual(1, corpus[0][0])
+
 
 if __name__ == '__main__':
     suite = unittest.TestSuite()