<NME> sampler.py
<BEF> from scannerpy.common import *
from scannerpy.protobufs import protobufs

DEFAULT_TASK_SIZE = 125

class Sampler:
    """
    Utility for specifying which frames of a video (or which rows of a table)
    to run a computation over.
    """

    def __init__(self, sc):
        self._sc = sc

    def All(self, input):
        def arg_builder():
            sampling_args = protobufs.SamplingArgs()
            sampling_args.sampling_function = "All"
            return sampling_args
        return self._sc.ops.Sample(
            col=self,
            args={'arg_builder': arg_builder,
                  'default': arg_builder()})

    def Stride(self, input, stride=None):
        def arg_builder(stride=stride):
            args = protobufs.StridedSamplerArgs()
            args.stride = stride
            sampling_args = protobufs.SamplingArgs()
            sampling_args.sampling_function = "Strided"
            sampling_args.sampling_args = args.SerializeToString()
            return sampling_args

        return self._sc.ops.Sample(
            col=self,
            args={'arg_builder': arg_builder,
                  'default': arg_builder() if stride else None})

    def Range(self, input, start=None, end=None):
        return self.ranges([(start, end)])

    def Ranges(self, input, intervals=None):
        return self.strided_ranges(intervals, 1)

    def Gather(self, input, rows=None):
        def arg_builder(rows=rows):
            args = protobufs.GatherSamplerArgs()
            args.rows[:] = rows
        return tasks

    def range(self, videos, start, end, item_size=1000, warmup_size=0):
        videos = self._convert_collection(videos)
        tasks = []
        for video in videos:
            task = self.strided_range(video, start, end, 1,
                                      item_size=item_size,
                                      warmup_size=warmup_size)
            tasks.append(task)
    def StridedRanges(self, input, intervals=None, stride=None):
        def arg_builder(intervals=intervals, stride=stride):
            args = protobufs.StridedRangeSamplerArgs()
            args.stride = stride
            for start, end in intervals:
                args.starts.append(start)
                args.ends.append(end)
            sampling_args = protobufs.SamplingArgs()
            sampling_args.sampling_function = "StridedRanges"
            sampling_args.sampling_args = args.SerializeToString()
            return sampling_args

        return self._sc.ops.Sample(
            col=self,
            args={'arg_builder': arg_builder,
                  'default': arg_builder() if intervals and stride else None})

    def RepeatNull(self, input, spacing=None):
        def arg_builder(spacing=spacing):
            args = protobufs.SpaceNullSamplerArgs()
            args.spacing = spacing
            sampling_args = protobufs.SamplingArgs()
            sampling_args.sampling_function = "SpaceNull"
            sampling_args.sampling_args = args.SerializeToString()
            return sampling_args

        return sampling_args

    def strided_range(self, video, start, end, stride, item_size=1000,
                      warmup_size=0):
        if isinstance(video, list) or isinstance(video, Collection):
            raise ScannerException('Sampler.strided_range only takes a single video')
        if not isinstance(video, tuple):
            return sampling_args

        return self._sc.ops.Space(
            col=self,
            args={'arg_builder': arg_builder,
                  'default': arg_builder() if spacing else None})
        sample.sampling_function = "StridedRange"
        sampler_args = self._db.protobufs.StridedRangeSamplerArgs()
        sampler_args.stride = stride
        s = start
        while s < end:
            ws = max(0, s - warmup_size * stride)
            e = min(s + item_size * stride, end)
            sampler_args.warmup_starts.append(ws)
            sampler_args.starts.append(s)
            sampler_args.ends.append(e)
            s = e
        sample.sampling_args = sampler_args.SerializeToString()
        return task

<MSG> Add ranges to sampler api that takes list of intervals

<DFF> @@ -49,10 +49,14 @@ class Sampler:
         return tasks
 
     def range(self, videos, start, end, item_size=1000, warmup_size=0):
+        return self.ranges(videos, [(start, end)], item_size=item_size,
+                           warmup_size=warmup_size)
+
+    def ranges(self, videos, intervals, item_size=1000, warmup_size=0):
         videos = self._convert_collection(videos)
         tasks = []
         for video in videos:
-            task = self.strided_range(video, start, end, 1,
+            task = self.strided_ranges(video, intervals, 1,
                                       item_size=item_size,
                                       warmup_size=warmup_size)
             tasks.append(task)
@@ -86,6 +90,12 @@ or (input_table, output_table) pair')""")
 
     def strided_range(self, video, start, end, stride, item_size=1000,
                       warmup_size=0):
+        return self.strided_ranges(video, [(start, end)], stride,
+                                   item_size=item_size,
+                                   warmup_size=warmup_size)
+
+    def strided_ranges(self, video, intervals, stride, item_size=1000,
+                      warmup_size=0):
         if isinstance(video, list) or isinstance(video, Collection):
             raise ScannerException('Sampler.strided_range only takes a single video')
         if not isinstance(video, tuple):
@@ -104,13 +114,14 @@ or (input_table, output_table) pair')""")
         sample.sampling_function = "StridedRange"
         sampler_args = self._db.protobufs.StridedRangeSamplerArgs()
         sampler_args.stride = stride
-        s = start
-        while s < end:
-            ws = max(0, s - warmup_size * stride)
-            e = min(s + item_size * stride, end)
-            sampler_args.warmup_starts.append(ws)
-            sampler_args.starts.append(s)
-            sampler_args.ends.append(e)
-            s = e
+        for start, end in intervals:
+            s = start
+            while s < end:
+                ws = max(0, s - warmup_size * stride)
+                e = min(s + item_size * stride, end)
+                sampler_args.warmup_starts.append(ws)
+                sampler_args.starts.append(s)
+                sampler_args.ends.append(e)
+                s = e
         sample.sampling_args = sampler_args.SerializeToString()
         return task

{'addition_count': 20, 'commit_subject': 'Add ranges to sampler api that takes list of intervals', 'deletion_count': 9, 'file_extension': '.py', 'lang': 'py', 'license': 'apache-2.0', 'repo_name': 'scanner-research/scanner'}