Repository: piskvorky/gensim

Issue:
Allow dictionary.save to take an open file handle.
It would improve composability if the `save` method of a `Dictionary` instance could take an open file handle instead of a filename. In particular, there are certain frameworks which only pass around file-like objects and do not use filenames, such as `luigi`. This change would increase compatibility from frameworks like these and provide room for other custom file-like objects to be used.



Discussion:
That would only work when saving into a single file (no `mmap`, no extra numpy/scipy.sparse array files being saved). Which is always the case with Dictionary, I think, so no problem.

Implementing should be easy -- a few changes around here https://github.com/piskvorky/gensim/blob/develop/gensim/utils.py#L246 = a special code path for when `fname` is a file handle. Let me know if you want to open a pull request, I'm +1 on this feature.


Resolving patch:
diff --git a/gensim/corpora/dictionary.py b/gensim/corpora/dictionary.py
--- a/gensim/corpora/dictionary.py
+++ b/gensim/corpora/dictionary.py
@@ -392,4 +392,5 @@ def from_corpus(corpus, id2word=None):
         logger.info("built %s from %i documents (total %i corpus positions)" %
                      (result, result.num_docs, result.num_pos))
         return result
+
 #endclass Dictionary
diff --git a/gensim/utils.py b/gensim/utils.py
--- a/gensim/utils.py
+++ b/gensim/utils.py
@@ -214,47 +214,58 @@ def load(cls, fname, mmap=None):
         subname = lambda suffix: fname + '.' + suffix + '.npy'
         obj = unpickle(fname)
         for attrib in getattr(obj, '__numpys', []):
-            logger.info("loading %s from %s with mmap=%s" % (attrib, subname(attrib), mmap))
+            logger.info("loading %s from %s with mmap=%s" % (
+                attrib, subname(attrib), mmap))
             setattr(obj, attrib, numpy.load(subname(attrib), mmap_mode=mmap))
+
         for attrib in getattr(obj, '__scipys', []):
-            logger.info("loading %s from %s with mmap=%s" % (attrib, subname(attrib), mmap))
+            logger.info("loading %s from %s with mmap=%s" % (
+                attrib, subname(attrib), mmap))
             sparse = unpickle(subname(attrib))
             sparse.data = numpy.load(subname(attrib) + '.data.npy', mmap_mode=mmap)
             sparse.indptr = numpy.load(subname(attrib) + '.indptr.npy', mmap_mode=mmap)
             sparse.indices = numpy.load(subname(attrib) + '.indices.npy', mmap_mode=mmap)
             setattr(obj, attrib, sparse)
+
         for attrib in getattr(obj, '__ignoreds', []):
             logger.info("setting ignored attribute %s to None" % (attrib))
             setattr(obj, attrib, None)
         return obj
 
-    def save(self, fname, separately=None, sep_limit=10 * 1024**2, ignore=frozenset()):
+    def _smart_save(self, fname, separately=None, sep_limit=10 * 1024**2,
+                    ignore=frozenset()):
         """
         Save the object to file (also see `load`).
 
-        If `separately` is None, automatically detect large numpy/scipy.sparse arrays
-        in the object being stored, and store them into separate files. This avoids
-        pickle memory errors and allows mmap'ing large arrays back on load efficiently.
+        If `separately` is None, automatically detect large
+        numpy/scipy.sparse arrays in the object being stored, and store
+        them into separate files. This avoids pickle memory errors and
+        allows mmap'ing large arrays back on load efficiently.
 
-        You can also set `separately` manually, in which case it must be a list of attribute
-        names to be stored in separate files. The automatic check is not performed in this case.
+        You can also set `separately` manually, in which case it must be
+        a list of attribute names to be stored in separate files. The
+        automatic check is not performed in this case.
 
-        `ignore` is a set of attribute names to *not* serialize (file handles, caches etc). On
-        subsequent load() these attributes will be set to None.
+        `ignore` is a set of attribute names to *not* serialize (file
+        handles, caches etc). On subsequent load() these attributes will
+        be set to None.
 
         """
-        logger.info("saving %s object under %s, separately %s" % (self.__class__.__name__, fname, separately))
+        logger.info(
+            "saving %s object under %s, separately %s" % (
+                self.__class__.__name__, fname, separately))
         subname = lambda suffix: fname + '.' + suffix + '.npy'
         tmp = {}
+        sparse_matrices = (scipy.sparse.csr_matrix, scipy.sparse.csc_matrix)
         if separately is None:
             separately = []
             for attrib, val in iteritems(self.__dict__):
                 if isinstance(val, numpy.ndarray) and val.size >= sep_limit:
                     separately.append(attrib)
-                elif isinstance(val, (scipy.sparse.csr_matrix, scipy.sparse.csc_matrix)) and val.nnz >= sep_limit:
+                elif isinstance(val, sparse_matrices) and val.nnz >= sep_limit:
                     separately.append(attrib)
 
-        # whatever's in `separately` or `ignore` at this point won't get pickled anymore
+        # whatever's in `separately` or `ignore` at this point won't get pickled
         for attrib in separately + list(ignore):
             if hasattr(self, attrib):
                 tmp[attrib] = getattr(self, attrib)
@@ -265,16 +276,19 @@ def save(self, fname, separately=None, sep_limit=10 * 1024**2, ignore=frozenset(
             for attrib, val in iteritems(tmp):
                 if isinstance(val, numpy.ndarray) and attrib not in ignore:
                     numpys.append(attrib)
-                    logger.info("storing numpy array '%s' to %s" % (attrib, subname(attrib)))
+                    logger.info("storing numpy array '%s' to %s" % (
+                        attrib, subname(attrib)))
                     numpy.save(subname(attrib), numpy.ascontiguousarray(val))
-                elif isinstance(val, (scipy.sparse.csr_matrix, scipy.sparse.csc_matrix)) and attrib not in ignore:
+                elif isinstance(val, sparse_matrices) and attrib not in ignore:
                     scipys.append(attrib)
-                    logger.info("storing scipy.sparse array '%s' under %s" % (attrib, subname(attrib)))
+                    logger.info("storing scipy.sparse array '%s' under %s" % (
+                        attrib, subname(attrib)))
                     numpy.save(subname(attrib) + '.data.npy', val.data)
                     numpy.save(subname(attrib) + '.indptr.npy', val.indptr)
                     numpy.save(subname(attrib) + '.indices.npy', val.indices)
                     data, indptr, indices = val.data, val.indptr, val.indices
                     val.data, val.indptr, val.indices = None, None, None
+
                     try:
                         pickle(val, subname(attrib)) # store array-less object
                     finally:
@@ -282,6 +296,7 @@ def save(self, fname, separately=None, sep_limit=10 * 1024**2, ignore=frozenset(
                 else:
                     logger.info("not storing attribute %s" % (attrib))
                     ignoreds.append(attrib)
+
             self.__dict__['__numpys'] = numpys
             self.__dict__['__scipys'] = scipys
             self.__dict__['__ignoreds'] = ignoreds
@@ -290,11 +305,41 @@ def save(self, fname, separately=None, sep_limit=10 * 1024**2, ignore=frozenset(
             # restore the attributes
             for attrib, val in iteritems(tmp):
                 setattr(self, attrib, val)
+
+    def save(self, fname_or_handle, separately=None, sep_limit=10 * 1024**2,
+             ignore=frozenset()):
+        """
+        Save the object to file (also see `load`).
+
+        `fname_or_handle` is either a string specifying the file name to
+        save to, or an open file-like object which can be written to. If
+        the object is a file handle, no special array handling will be
+        performed; all attributes will be saved to the same file.
+
+        If `separately` is None, automatically detect large
+        numpy/scipy.sparse arrays in the object being stored, and store
+        them into separate files. This avoids pickle memory errors and
+        allows mmap'ing large arrays back on load efficiently.
+
+        You can also set `separately` manually, in which case it must be
+        a list of attribute names to be stored in separate files. The
+        automatic check is not performed in this case.
+
+        `ignore` is a set of attribute names to *not* serialize (file
+        handles, caches etc). On subsequent load() these attributes will
+        be set to None.
+
+        """
+        try:
+            _pickle.dump(self, fname_or_handle, protocol=_pickle.HIGHEST_PROTOCOL)
+            logger.info("saved %s object" % self.__class__.__name__)
+        except TypeError:  # `fname_or_handle` does not have write attribute
+            self._smart_save(fname_or_handle, separately, sep_limit, ignore)
 #endclass SaveLoad
 
 
 def identity(p):
-    """Identity fnc, for flows that don't accept lambda (picking etc)."""
+    """Identity fnc, for flows that don't accept lambda (pickling etc)."""
     return p
 
 
@@ -715,7 +760,7 @@ def smart_open(fname, mode='rb'):
     return open(fname, mode)
 
 
-def pickle(obj, fname, protocol=-1):
+def pickle(obj, fname, protocol=_pickle.HIGHEST_PROTOCOL):
     """Pickle object `obj` to file `fname`."""
     with smart_open(fname, 'wb') as fout: # 'b' for binary, needed on Windows
         _pickle.dump(obj, fout, protocol=protocol)


Test patch: