X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/19c43cebdfa1b701b031577f62c2a7b060dc8b0e..f34c10a5b22e3f9941f2cfe60f6f9d7a81319b6f:/sdk/python/arvados/_ranges.py

diff --git a/sdk/python/arvados/_ranges.py b/sdk/python/arvados/_ranges.py
index d5ff6ed1b5..83437b2adb 100644
--- a/sdk/python/arvados/_ranges.py
+++ b/sdk/python/arvados/_ranges.py
@@ -21,7 +21,7 @@ class Range(object):
                 self.range_size == other.range_size and
                 self.segment_offset == other.segment_offset)
 
-def first_block(data_locators, range_start, range_size):
+def first_block(data_locators, range_start):
     block_start = 0L
 
     # range_start/block_start is the inclusive lower bound
@@ -68,7 +68,7 @@ class LocatorAndRange(object):
     def __repr__(self):
         return "LocatorAndRange(%r, %r, %r, %r)" % (self.locator, self.block_size, self.segment_offset, self.segment_size)
 
-def locators_and_ranges(data_locators, range_start, range_size):
+def locators_and_ranges(data_locators, range_start, range_size, limit=None):
     """Get blocks that are covered by a range.
 
     Returns a list of LocatorAndRange objects.
@@ -82,19 +82,24 @@ def locators_and_ranges(data_locators, range_start, range_size):
     :range_size:
       size of range
 
+    :limit:
+      Maximum segments to return, default None (unlimited).  Will truncate the
+      result if there are more segments needed to cover the range than the
+      limit.
+
     """
     if range_size == 0:
         return []
     resp = []
     range_end = range_start + range_size
 
-    i = first_block(data_locators, range_start, range_size)
+    i = first_block(data_locators, range_start)
     if i is None:
         return []
 
     # We should always start at the first segment due to the binary
     # search.
-    while i < len(data_locators):
+    while i < len(data_locators) and len(resp) != limit:
         dl = data_locators[i]
         block_start = dl.range_start
         block_size = dl.range_size
@@ -163,7 +168,7 @@ def replace_range(data_locators, new_range_start, new_range_size, new_locator, n
             data_locators.append(Range(new_locator, new_range_start, new_range_size, new_segment_offset))
         return
 
-    i = first_block(data_locators, new_range_start, new_range_size)
+    i = first_block(data_locators, new_range_start)
     if i is None:
         return