Merge pull request #146 from jimrybarski/132-range-for-select

resolves #132: The `select` method now allows you to specify a range …

functional_tests/FYLM141111001.py

@@ -7,6 +7,7 @@ from nd2reader import Nd2
 import numpy as np
 from datetime import datetime
 import unittest
+import time
 
 
 class FYLM141111Tests(unittest.TestCase):
@@ -171,3 +172,44 @@ class FYLM141111Tests(unittest.TestCase):
             n = image.index
             if n > 100:
                 break
+
+    def test_select_start(self):
+        count = 0
+        for _ in self.nd2.select(channels='GFP', start=29000):
+            count += 1
+        self.assertEqual(127, count)
+
+    def test_select_stop(self):
+        count = 0
+        for _ in self.nd2.select(channels='GFP', stop=20):
+            count += 1
+        self.assertEqual(count, 3)
+
+    def test_select_start_stop(self):
+        count = 0
+        for _ in self.nd2.select(channels='GFP', start=10, stop=20):
+            count += 1
+        self.assertEqual(count, 1)
+
+    def test_select_start_stop_brightfield(self):
+        count = 0
+        for _ in self.nd2.select(channels='', start=10, stop=20):
+            count += 1
+        self.assertEqual(count, 5)
+
+    def test_select_faster(self):
+        select_count = 0
+        select_start = time.time()
+        for i in self.nd2.select(channels='GFP', start=10, stop=50):
+            if i is not None and i.channel == 'GFP':
+                select_count += 1
+        select_duration = time.time() - select_start
+
+        direct_count = 0
+        direct_start = time.time()
+        for i in self.nd2[10:50]:
+            if i is not None and i.channel == 'GFP':
+                direct_count += 1
+        direct_duration = time.time() - direct_start
+        self.assertEqual(select_count, direct_count)
+        self.assertGreater(direct_duration, select_duration)

functional_tests/monocycle.py

@@ -81,9 +81,11 @@ class Monocycle2Tests(unittest.TestCase):
     def tearDown(self):
         self.nd2.close()
 
+    @unittest.skip('missing file')
     def test_pixel_size(self):
         self.assertGreater(self.nd2.pixel_microns, 0.0)
 
+    @unittest.skip('missing file')
     def test_select(self):
         manual_images = []
         for _, image in zip(range(20), self.nd2):
@@ -104,6 +106,7 @@ class Monocycle2Tests(unittest.TestCase):
             self.assertEqual(a.field_of_view, b.field_of_view)
             self.assertEqual(a.channel, b.channel)
 
+    @unittest.skip('missing file')
     def test_select_order_all(self):
         # If we select every possible image using select(), we should just get every image in order
         n = 0
@@ -122,6 +125,7 @@ class Monocycle2Tests(unittest.TestCase):
                 # If there's a problem, we'll have seen it by now.
                 break
 
+    @unittest.skip('missing file')
     def test_select_order_subset(self):
         # Test that images are always yielded in increasing order. This guarantees that no matter what subset of images
         # we're filtering, we still get them in the chronological order they were acquired

functional_tests/single.py

@@ -62,3 +62,7 @@ class SingleTests(unittest.TestCase):
     def test_iteration_backwards(self):
         images = [image for image in self.nd2[::-1]]
         self.assertEqual(len(images), 1)
+
+    def test_select_bounds_wrong(self):
+        images = [i for i in self.nd2.select(start=0, stop=12481247)]
+        self.assertEqual(len(images), 1)

nd2reader/main.py

@@ -61,7 +61,7 @@ class Nd2(object):
             return self._slice(item.start, item.stop, item.step)
         raise IndexError
 
-    def select(self, fields_of_view=None, channels=None, z_levels=None):
+    def select(self, fields_of_view=None, channels=None, z_levels=None, start=0, stop=None):
         """
         Iterates over images matching the given criteria. This can be 2-10 times faster than manually iterating over
         the Nd2 and checking the attributes of each image, as this method skips disk reads for any images that don't
@@ -70,13 +70,17 @@ class Nd2(object):
         :type fields_of_view:   int or tuple or list
         :type channels:     str or tuple or list
         :type z_levels:     int or tuple or list
+        :type start:        int
+        :type stop:         int
 
         """
         fields_of_view = self._to_tuple(fields_of_view, self.fields_of_view)
         channels = self._to_tuple(channels, self.channels)
         z_levels = self._to_tuple(z_levels, self.z_levels)
 
-        for frame in range(len(self)):
+        # By default, we stop after the last image. Otherwise we make sure the user-provided value is valid
+        stop = len(self) if stop is None else max(0, min(stop, len(self)))
+        for frame in range(start, stop):
             field_of_view, channel, z_level = self._parser.driver.calculate_image_properties(frame)
             if field_of_view in fields_of_view and channel in channels and z_level in z_levels:
                 image = self._parser.driver.get_image(frame)