util: add limit to amount filechunkiter will read
authorVadim Gelfer <vadim.gelfer@gmail.com>
Tue, 20 Jun 2006 15:13:17 -0700
changeset 2462 d610bcfd66a8
parent 2461 ff9ee834e3b6
child 2463 6ab016edd5c4
util: add limit to amount filechunkiter will read
mercurial/util.py
--- a/mercurial/util.py	Tue Jun 20 09:11:41 2006 -0700
+++ b/mercurial/util.py	Tue Jun 20 15:13:17 2006 -0700
@@ -822,16 +822,22 @@
         s, self.buf = self.buf[:l], buffer(self.buf, l)
         return s
 
-def filechunkiter(f, size = 65536):
-    """Create a generator that produces all the data in the file size
-    (default 65536) bytes at a time.  Chunks may be less than size
-    bytes if the chunk is the last chunk in the file, or the file is a
-    socket or some other type of file that sometimes reads less data
-    than is requested."""
-    s = f.read(size)
-    while len(s) > 0:
+def filechunkiter(f, size=65536, limit=None):
+    """Create a generator that produces the data in the file size
+    (default 65536) bytes at a time, up to optional limit (default is
+    to read all data).  Chunks may be less than size bytes if the
+    chunk is the last chunk in the file, or the file is a socket or
+    some other type of file that sometimes reads less data than is
+    requested."""
+    assert size >= 0
+    assert limit is None or limit >= 0
+    while True:
+        if limit is None: nbytes = size
+        else: nbytes = min(limit, size)
+        s = nbytes and f.read(nbytes)
+        if not s: break
+        if limit: limit -= len(s)
         yield s
-        s = f.read(size)
 
 def makedate():
     lt = time.localtime()