Repository: DataBiosphere/toil

Issue:
writeGlobalFileStream should return FileID object with .size
Issue #1014 woks on adding attributes to the return value of writeGlobalFile. This issue concerns setting up the same for writeGlobalFileStream.

┆Issue is synchronized with this [JIRA Story](https://ucsc-cgl.atlassian.net/browse/TOIL-106)
┆Issue Number: TOIL-106



Discussion:
This is causing me problems now. The documentation claims that writeGlobalFileStream gives the file ID of the file being written, but really it just gives an str. Even when passing it along to successor jobs, it's still not a file ID with a size member, as the documentation suggests it should be. So it's possible to get files in your file store with undeterminable sizes, which in turn interferes with jobs that process files and make sure to reserve the right amount of disk to process the files they are dealing with.

Resolving patch:
diff --git a/src/toil/fileStore.py b/src/toil/fileStore.py
--- a/src/toil/fileStore.py
+++ b/src/toil/fileStore.py
@@ -180,6 +180,7 @@ def writeGlobalFile(self, localFileName, cleanup=False):
         """
         raise NotImplementedError()
 
+    @contextmanager
     def writeGlobalFileStream(self, cleanup=False):
         """
         Similar to writeGlobalFile, but allows the writing of a stream to the job store.
@@ -188,10 +189,28 @@ def writeGlobalFileStream(self, cleanup=False):
         :param bool cleanup: is as in :func:`toil.fileStore.FileStore.writeGlobalFile`.
         :return: A context manager yielding a tuple of
                   1) a file handle which can be written to and
-                  2) the ID of the resulting file in the job store.
+                  2) the toil.fileStore.FileID of the resulting file in the job store.
         """
+        
         # TODO: Make this work with FileID
-        return self.jobStore.writeFileStream(None if not cleanup else self.jobGraph.jobStoreID)
+        with self.jobStore.writeFileStream(None if not cleanup else self.jobGraph.jobStoreID) as (backingStream, fileStoreID):
+            # We have a string version of the file ID, and the backing stream.
+            # We need to yield a stream the caller can write to, and a FileID
+            # that accurately reflects the size of the data written to the
+            # stream. We assume the stream is not seekable.
+            
+            # Make and keep a reference to the file ID, which is currently empty
+            fileID = FileID(fileStoreID, 0)
+            
+            # Wrap the stream to increment the file ID's size for each byte written
+            wrappedStream = WriteWatchingStream(backingStream)
+            
+            # When the stream is written to, count the bytes
+            def handle(numBytes):
+                fileID.size += numBytes 
+            wrappedStream.onWrite(handle)
+            
+            yield wrappedStream, fileID
 
     @abstractmethod
     def readGlobalFile(self, fileStoreID, userPath=None, cache=True, mutable=False, symlink=False):
@@ -1847,6 +1866,65 @@ def __init__(self, fileStoreID, size):
     @classmethod
     def forPath(cls, fileStoreID, filePath):
         return cls(fileStoreID, os.stat(filePath).st_size)
+        
+class WriteWatchingStream(object):
+    """
+    A stream wrapping class that calls any functions passed to onWrite() with the number of bytes written for every write.
+    
+    Not seekable.
+    """
+    
+    def __init__(self, backingStream):
+        """
+        Wrap the given backing stream.
+        """
+        
+        self.backingStream = backingStream
+        # We have no write listeners yet
+        self.writeListeners = []
+        
+    def onWrite(self, listener):
+        """
+        Call the given listener with the number of bytes written on every write.
+        """
+        
+        self.writeListeners.append(listener)
+        
+    # Implement the file API from https://docs.python.org/2.4/lib/bltin-file-objects.html
+        
+    def write(self, data):
+        """
+        Write the given data to the file.
+        """
+        
+        # Do the write
+        self.backingStream.write(data)
+        
+        for listener in self.writeListeners:
+            # Send out notifications
+            listener(len(data))
+            
+    def writelines(self, datas):
+        """
+        Write each string from the given iterable, without newlines.
+        """
+        
+        for data in datas:
+            self.write(data)
+            
+    def flush(self):
+        """
+        Flush the backing stream.
+        """
+        
+        self.backingStream.flush()
+        
+    def close(self):
+        """
+        Close the backing stream.
+        """
+        
+        self.backingStream.close()
 
 
 def shutdownFileStore(workflowDir, workflowID):


Test patch:
diff --git a/src/toil/test/src/jobFileStoreTest.py b/src/toil/test/src/jobFileStoreTest.py
--- a/src/toil/test/src/jobFileStoreTest.py
+++ b/src/toil/test/src/jobFileStoreTest.py
@@ -23,6 +23,7 @@
 
 from toil.common import Toil
 from toil.job import Job
+from toil.fileStore import FileID
 from toil.test import ToilTest, slow
 
 PREFIX_LENGTH=200
@@ -134,12 +135,22 @@ def fileTestJob(job, inputFileStoreIDs, testStrings, chainLength):
             with open(tempFile, 'w') as fH:
                 fH.write(testString)
             #Write a local copy of the file using the local file
-            outputFileStoreIds.append(job.fileStore.writeGlobalFile(tempFile))
+            fileStoreID = job.fileStore.writeGlobalFile(tempFile)
+            
+            # Make sure it returned a valid and correct FileID with the right size
+            assert isinstance(fileStoreID, FileID)
+            assert fileStoreID.size == len(testString.encode('utf-8'))
+            
+            outputFileStoreIds.append(fileStoreID)
         else:
             #Use the writeGlobalFileStream method to write the file
             with job.fileStore.writeGlobalFileStream() as (fH, fileStoreID):
                 fH.write(testString.encode('utf-8'))
                 outputFileStoreIds.append(fileStoreID)
+                
+            #Make sure it returned a valid and correct FileID with the right size
+            assert isinstance(fileStoreID, FileID)
+            assert fileStoreID.size == len(testString.encode('utf-8'))
 
     if chainLength > 0:
         #Make a child that will read these files and check it gets the same results