From e141b0504cbb616b89ec516983a9e8339b5dd956 Mon Sep 17 00:00:00 2001
From: Jim Rybarski <jim@rybarski.com>
Date: Fri, 4 Dec 2015 00:15:43 -0600
Subject: [PATCH] #114: wrote fast filter that works, but we need way more
 testing for corner cases

---
 Makefile                          |  4 ++--
 functional_tests/FYLM141111001.py | 14 ++++++++++++++
 nd2reader/interface.py            | 25 +++++++++++++++++++++++++
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/Makefile b/Makefile
index 3ddd695..b2301d6 100644
--- a/Makefile
+++ b/Makefile
@@ -17,7 +17,7 @@ build:
 	docker build -t jimrybarski/nd2reader .
 
 shell:
-	xhost local:root; docker run --rm -v ~/nd2s:/var/nd2s -v /tmp/.X11-unix:/tmp/.X11-unix -e DISPLAY=unix$(DISPLAY) -it jimrybarski/nd2reader bash
+	xhost local:root; docker run --rm -v $(CURDIR):/opt/nd2reader -v ~/nd2s:/var/nd2s -v /tmp/.X11-unix:/tmp/.X11-unix -e DISPLAY=unix$(DISPLAY) -it jimrybarski/nd2reader bash
 
 py2:
 	xhost local:root; docker run --rm -v $(CURDIR):/opt/nd2reader -v ~/nd2s:/var/nd2s -v /tmp/.X11-unix:/tmp/.X11-unix -e DISPLAY=unix$(DISPLAY) -it jimrybarski/nd2reader python2.7
@@ -30,7 +30,7 @@ test:	build
 	docker run --rm -v $(CURDIR):/opt/nd2reader -it jimrybarski/nd2reader python2.7 test.py
 
 ftest:	build
-	docker run --rm -v $(CURDIR):/opt/nd2reader -v ~/nd2s:/var/nd2s -it jimrybarski/nd2reader python3.4 /opt/nd2reader/ftest.py
+	xhost local:root; docker run --rm -v /tmp/.X11-unix:/tmp/.X11-unix -e DISPLAY=unix$(DISPLAY) -v $(CURDIR):/opt/nd2reader -v ~/nd2s:/var/nd2s -it jimrybarski/nd2reader python3.4 /opt/nd2reader/ftest.py
 	docker run --rm -v $(CURDIR):/opt/nd2reader -v ~/nd2s:/var/nd2s -it jimrybarski/nd2reader python2.7 /opt/nd2reader/ftest.py
 	
 publish:
diff --git a/functional_tests/FYLM141111001.py b/functional_tests/FYLM141111001.py
index 9d1d6df..ad1ec59 100644
--- a/functional_tests/FYLM141111001.py
+++ b/functional_tests/FYLM141111001.py
@@ -4,6 +4,7 @@ run them unless you're Jim Rybarski.
 
 """
 from nd2reader import Nd2
+import numpy as np
 from datetime import datetime
 import unittest
 
@@ -109,3 +110,16 @@ class FunctionalTests(unittest.TestCase):
     def test_get_image_by_attribute_none(self):
         image = self.nd2.get_image(4, 0, "GFP", 0)
         self.assertIsNone(image)
+
+    def test_fast_filter(self):
+        manual_images = []
+        for _, image in zip(range(200), self.nd2):
+            if image is not None and image.channel == 'GFP':
+                manual_images.append(image)
+        filter_images = []
+        for image in self.nd2.filter(channels=['GFP']):
+            filter_images.append(image)
+            if len(filter_images) == len(manual_images):
+                break
+        for a, b in zip(manual_images, filter_images):
+            self.assertTrue(np.array_equal(a, b))
diff --git a/nd2reader/interface.py b/nd2reader/interface.py
index 32a4887..efda766 100644
--- a/nd2reader/interface.py
+++ b/nd2reader/interface.py
@@ -2,6 +2,7 @@
 
 from nd2reader.parser import get_parser
 from nd2reader.version import get_version
+import six
 
 
 class Nd2(object):
@@ -181,6 +182,29 @@ class Nd2(object):
                                                            self.height,
                                                            self.width)
 
+    def filter(self, fields_of_view=None, channels=None, z_levels=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 will not read from disk until a valid image is
+        found.
+
+        """
+        fields_of_view = self._to_list(fields_of_view, self.fields_of_view)
+        channels = self._to_list(channels, self.channels)
+        z_levels = self._to_list(z_levels, self.z_levels)
+
+        for frame in self.frames:
+            for fov in fields_of_view:
+                for z in z_levels:
+                    for c in channels:
+                        image = self.get_image(frame, fov, c, z)
+                        if image is not None:
+                            yield image
+
+    def _to_list(self, value, default):
+        value = default if value is None else value
+        return [value] if isinstance(value, int) or isinstance(value, six.string_types) else value
+
     def close(self):
         """
         Closes the file handle to the image. This actually sometimes will prevent problems so it's good to do this or
@@ -188,3 +212,4 @@ class Nd2(object):
 
         """
         self._fh.close()
+1
\ No newline at end of file