From a40b243abb0685a6d873ee554807840154b5c680 Mon Sep 17 00:00:00 2001
From: Jim Rybarski <jim@rybarski.com>
Date: Sun, 31 May 2015 05:31:06 +0000
Subject: [PATCH] resolves #40: adds slicing semantics

---
 nd2reader/__init__.py | 50 +++++++++++++++++++++++++++++++++++++------
 1 file changed, 43 insertions(+), 7 deletions(-)

diff --git a/nd2reader/__init__.py b/nd2reader/__init__.py
index 754c6ef..7f9cf1e 100644
--- a/nd2reader/__init__.py
+++ b/nd2reader/__init__.py
@@ -34,19 +34,55 @@ class Nd2(Nd2Parser):
         """
         return self._image_count * self._channel_count
 
-    def __getitem__(self, frame_number):
-        if isinstance(frame_number, int):
+    def __getitem__(self, item):
+        """
+        Allows slicing ND2s.
+
+        >>> nd2 = Nd2("my_images.nd2")
+        >>> image = nd2[16]  # gets 17th frame
+        >>> for image in nd2[100:200]:  # iterate over the 100th to 200th images
+        >>>     do_something(image.data)
+        >>> for image in nd2[::-1]:  # iterate backwards
+        >>>     do_something(image.data)
+        >>> for image in nd2[37:422:17]:  # do something super weird if you really want to
+        >>>     do_something(image.data)
+
+        :type item: int or slice
+        :rtype: nd2reader.model.Image() or generator
+
+        """
+        if isinstance(item, int):
             try:
-                channel_offset = frame_number % len(self.channels)
-                fov = self._calculate_field_of_view(frame_number)
-                channel = self._calculate_channel(frame_number)
-                z_level = self._calculate_z_level(frame_number)
-                timestamp, raw_image_data = self._get_raw_image_data(frame_number, channel_offset)
+                channel_offset = item % len(self.channels)
+                fov = self._calculate_field_of_view(item)
+                channel = self._calculate_channel(item)
+                z_level = self._calculate_z_level(item)
+                timestamp, raw_image_data = self._get_raw_image_data(item, channel_offset)
                 image = Image(timestamp, raw_image_data, fov, channel, z_level, self.height, self.width)
             except (TypeError, ValueError):
                 return None
             else:
                 return image
+        elif isinstance(item, slice):
+            return self._slice(item.start, item.stop, item.step)
+        raise IndexError
+
+    def _slice(self, start, stop, step):
+        """
+        Allows for iteration over a selection of the entire dataset.
+
+        :type start: int
+        :type stop: int
+        :type step: int
+        :rtype: nd2reader.model.Image() or None
+
+        """
+        start = start if start is not None else 0
+        step = step if step is not None else 1
+        stop = stop if stop is not None else len(self)
+        # This weird thing with the step allows you to iterate backwards over the images
+        for i in range(start, stop)[::step]:
+            yield self[i]
 
     @property
     def image_sets(self):