From 731a05214acc82ecbaf926552477fa530527259d Mon Sep 17 00:00:00 2001
From: Ruben Verweij <verweij@physics.leidenuniv.nl>
Date: Thu, 4 Jan 2018 14:08:13 +0100
Subject: [PATCH] Fix #4

---
 nd2reader/raw_metadata.py | 25 +++++++++++++++++++++++--
 nd2reader/reader.py       |  6 ++++--
 tests/test_reader.py      |  2 +-
 3 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/nd2reader/raw_metadata.py b/nd2reader/raw_metadata.py
index 67ae7e1..b0d3485 100644
--- a/nd2reader/raw_metadata.py
+++ b/nd2reader/raw_metadata.py
@@ -377,12 +377,14 @@ class RawMetadata(object):
         for loop in loops:
             # duration of this loop
             duration = get_from_dict_if_exists('dDuration', loop) or 0
+            interval = self._determine_sampling_interval(duration, loop)
+
+            # if duration is not saved, infer it
+            duration = self.get_duration_from_interval_and_loops(duration, interval, loop)
 
             # uiLoopType == 6 is a stimulation loop
             is_stimulation = get_from_dict_if_exists('uiLoopType', loop) == 6
 
-            interval = self._determine_sampling_interval(duration, loop)
-
             parsed_loop = {
                 'start': time_offset,
                 'duration': duration,
@@ -397,6 +399,25 @@ class RawMetadata(object):
 
         return parsed_loops
 
+    def get_duration_from_interval_and_loops(self, duration, interval, loop):
+        """Infers the duration of the loop from the number of measurements and the interval
+
+        Args:
+            duration: loop duration in milliseconds
+            duration: measurement interval in milliseconds
+            loop: loop dictionary
+
+        Returns:
+            float: the loop duration in milliseconds
+
+        """
+        if duration == 0 and interval > 0:
+            number_of_loops = get_from_dict_if_exists('uiCount', loop)
+            number_of_loops = number_of_loops if number_of_loops is not None and number_of_loops > 0 else 1
+            duration = interval * number_of_loops
+
+        return duration
+
     @staticmethod
     def _determine_sampling_interval(duration, loop):
         """Determines the loop sampling interval in milliseconds
diff --git a/nd2reader/reader.py b/nd2reader/reader.py
index 30c3c1e..350d6f5 100644
--- a/nd2reader/reader.py
+++ b/nd2reader/reader.py
@@ -196,6 +196,8 @@ class ND2Reader(FramesSequenceND):
                 (timesteps, np.arange(current_time, current_time + loop['duration'], loop['sampling_interval'])))
             current_time += loop['duration']
 
-        # if experiment did not finish, number of timesteps is wrong. Take correct amount of leading timesteps.
-        self._timesteps = timesteps[:self.metadata['num_frames']]
+        if len(timesteps) > 0:
+            # if experiment did not finish, number of timesteps is wrong. Take correct amount of leading timesteps.
+            self._timesteps = timesteps[:self.metadata['num_frames']]
+
         return self._timesteps
diff --git a/tests/test_reader.py b/tests/test_reader.py
index 8c78249..71e607f 100644
--- a/tests/test_reader.py
+++ b/tests/test_reader.py
@@ -40,4 +40,4 @@ class TestReader(unittest.TestCase):
         with ArtificialND2('test_data/test_nd2_reader.nd2') as _:
             with ND2Reader('test_data/test_nd2_reader.nd2') as reader:
                 timesteps = reader.timesteps
-                self.assertEquals(list(timesteps), [])
+                self.assertEquals(timesteps, None)