From eeb4bd106328fb916822c10891945dbfc6ed95ec Mon Sep 17 00:00:00 2001
From: Ruben Verweij <ruben@lighthacking.nl>
Date: Mon, 4 May 2020 16:33:45 +0200
Subject: [PATCH] Fix issue #34

---
 nd2reader/raw_metadata.py  | 41 ++++++++++++++++++++++++++++----------
 setup.py                   |  1 -
 tests/test_raw_metadata.py | 11 ++++++++++
 tests/test_version.py      |  8 ++------
 4 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/nd2reader/raw_metadata.py b/nd2reader/raw_metadata.py
index 0d77ddc..1152219 100644
--- a/nd2reader/raw_metadata.py
+++ b/nd2reader/raw_metadata.py
@@ -158,18 +158,34 @@ class RawMetadata(object):
     def _parse_z_levels(self):
         """The different levels in the Z-plane.
 
+        If they are not available from the _parse_dimension function AND there
+        is NO 'Dimensions: ' textinfo item in the file, we return a range with
+        the length of z_coordinates if available, otherwise an empty list.
+
         Returns:
             list: the z levels, just a sequence from 0 to n.
         """
-        z_levels = self._parse_dimension(r""".*?Z\((\d+)\).*?""")
-        if 0 != len(z_levels):
-            z_levels = parse_if_not_none(self.z_data, self._parse_z_coordinates)
-            if z_levels is None:
-                z_levels = []
-            else:
-                z_levels = range(len(z_levels))
-                warnings.warn("Z-levels details missing in metadata. Using Z-coordinates instead.")
-        return z_levels
+        # get the dimension text to check if we should apply the fallback or not
+        dimension_text = self._parse_dimension_text()
+
+        # this returns range(len(z_levels))
+        z_levels = self._parse_dimension(r""".*?Z\((\d+)\).*?""", dimension_text)
+
+        if len(z_levels) > 0 or len(dimension_text) > 0:
+            # Either we have found the z_levels (first condition) so return, or
+            # don't fallback, because Z is apparently not in Dimensions, so
+            # there should be no z_levels
+            return z_levels
+
+        # Not available from dimension, get from z_coordinates
+        z_levels = parse_if_not_none(self.z_data, self._parse_z_coordinates)
+
+        if z_levels is None:
+            # No z coordinates, return empty list
+            return []
+
+        warnings.warn("Z-levels details missing in metadata. Using Z-coordinates instead.")
+        return range(len(z_levels))
 
     def _parse_z_coordinates(self):
         """The coordinate in micron for all z planes.
@@ -201,15 +217,18 @@ class RawMetadata(object):
 
         return dimension_text
 
-    def _parse_dimension(self, pattern):
-        dimension_text = self._parse_dimension_text()
+    def _parse_dimension(self, pattern, dimension_text=None):
+        dimension_text = self._parse_dimension_text() if dimension_text is None else dimension_text
         if dimension_text is None:
             return []
+
         if six.PY3:
             dimension_text = dimension_text.decode("utf8")
+
         match = re.match(pattern, dimension_text)
         if not match:
             return []
+
         count = int(match.group(1))
         return range(count)
 
diff --git a/setup.py b/setup.py
index 7688e05..0ade4f4 100644
--- a/setup.py
+++ b/setup.py
@@ -1,4 +1,3 @@
-from os import path
 from setuptools import setup
 from nd2reader import __version__ as VERSION
 
diff --git a/tests/test_raw_metadata.py b/tests/test_raw_metadata.py
index 914acd8..c21e0fe 100644
--- a/tests/test_raw_metadata.py
+++ b/tests/test_raw_metadata.py
@@ -30,6 +30,17 @@ class TestRawMetadata(unittest.TestCase):
         self.assertEqual(parse_dimension_text_line(line), six.b('Dimensions: T(443) x \xce\xbb(1)'))
         self.assertIsNone(parse_dimension_text_line(six.b('Dim: nothing')))
 
+    def test_parse_z_levels(self):
+        # smokescreen test to check if the fallback to z_coordinates is working
+        # for details, see RawMetadata._parse_z_levels()
+        dimension_text = self.metadata._parse_dimension_text()
+        z_levels = self.metadata._parse_dimension(r""".*?Z\((\d+)\).*?""", dimension_text)
+        z_coords = self.metadata._parse_z_coordinates()
+
+        self.assertEqual(len(dimension_text), 0)
+        self.assertEqual(len(z_levels), 0)
+        self.assertEqual(len(self.metadata._parse_z_levels()), len(z_coords))
+
     def test_dict(self):
         self.assertTrue(type(self.metadata.__dict__) is dict)
 
diff --git a/tests/test_version.py b/tests/test_version.py
index fedda9a..fedfeeb 100644
--- a/tests/test_version.py
+++ b/tests/test_version.py
@@ -1,12 +1,8 @@
-import nd2reader
 import unittest
-from setup import VERSION
+from nd2reader import __version__ as VERSION
 
 
 class TestVersion(unittest.TestCase):
     def test_module_version_type(self):
         # just make sure the version number exists and is the type we expect
-        self.assertEqual(type(nd2reader.__version__), str)
-
-    def test_versions_in_sync(self):
-        self.assertEqual(nd2reader.__version__, VERSION)
+        self.assertEqual(type(VERSION), str)