Added files.copyfileobj() as an alternative to shutil.copyfileobj() that supports a limit

Author: hannes-ucsc

src/bd2k/util/files.py

@@ -28,8 +28,47 @@ def rm_f( path ):
             raise
 
 
-if False:
+def copyfileobj( src, dst, limit=None, bufsize=1024 * 1024 ):
+    """
+    Copy the contents of one file object to another file object. If limit is given, stop after at
+    most limit bytes were copied. The copying will begin at the current file pointer of each file
+    object.
+
+    :param src: the file object to copy from
+
+    :param dst: the file object to copy to
+
+    :param limit: the maximum number of bytes to copy or None if all remaining bytes in src
+           should be copied
+
+    :param bufsize: the size of the intermediate copy buffer. No more than that many bytes will
+           ever be read from src or written to dst at any one time.
 
+    :return: None if limit is None, otherwise the difference between limit and the number of
+    bytes actually copied. This will be > 0 if and only if the source file hit EOF before limit
+    number of bytes could be read.
+
+    >>> import tempfile
+    >>> with open('/dev/urandom') as f1:
+    ...     with tempfile.TemporaryFile() as f2:
+    ...         copyfileobj(f1,f2,limit=100)
+    ...         f2.seek(60)
+    ...         with tempfile.TemporaryFile() as f3:
+    ...             copyfileobj(f2,f3), f2.tell(), f3.tell()
+    (None, 100, 40)
+    """
+    while limit is None or limit > 0:
+        buf = src.read( bufsize if limit is None or bufsize < limit else limit )
+        if buf:
+            if limit is not None:
+                limit -= len( buf )
+                assert limit >= 0
+            dst.write( buf )
+        else:
+            return limit
+
+
+if False:
     # These are not needed for Python 2.7 as Python's builtin file object's read() and write()
     # method are greedy. For Python 3.x these may be useful.