Merge pull request #39 from ggirelli/ggirelli-patch-nd2reader-from-file

Update to reader.py for issue #35

nd2reader/reader.py

@@ -11,18 +11,28 @@ class ND2Reader(FramesSequenceND):
     This is the main class: use this to process your .nd2 files.
     """
 
+    _fh = None
     class_priority = 12
 
-    def __init__(self, filename):
+    def __init__(self, fh):
+        """
+        Arguments:
+            fh {str} -- absolute path to .nd2 file
+            fh {IO} -- input buffer handler (opened with "rb" mode)
+        """
         super(ND2Reader, self).__init__()
 
-        if not filename.endswith(".nd2"):
-            raise InvalidFileType("The file %s you want to read with nd2reader does not have extension .nd2." % filename)
+        if isinstance(fh, str):
+            if not fh.endswith(".nd2"):
+                raise InvalidFileType(
+                    ("The file %s you want to read with nd2reader" % fh)
+                    + " does not have extension .nd2."
+                )
+            fh = open(fh, "rb")
 
-        self.filename = filename
+        self._fh = fh
+        self.filename = ""
 
-        # first use the parser to parse the file
-        self._fh = open(filename, "rb")
         self._parser = Parser(self._fh)
 
         # Setup metadata
@@ -42,7 +52,7 @@ class ND2Reader(FramesSequenceND):
         """Let PIMS open function use this reader for opening .nd2 files
 
         """
-        return {'nd2'} | super(ND2Reader, cls).class_exts()
+        return {"nd2"} | super(ND2Reader, cls).class_exts()
 
     def close(self):
         """Correctly close the file handle
@@ -119,22 +129,24 @@ class ND2Reader(FramesSequenceND):
     @property
     def frame_rate(self):
         """The (average) frame rate
-        
+
         Returns:
             float: the (average) frame rate in frames per second
         """
         total_duration = 0.0
 
-        for loop in self.metadata['experiment']['loops']:
-            total_duration += loop['duration']
+        for loop in self.metadata["experiment"]["loops"]:
+            total_duration += loop["duration"]
 
         if total_duration == 0:
             total_duration = self.timesteps[-1]
 
             if total_duration == 0:
-                raise ValueError('Total measurement duration could not be determined from loops')
+                raise ValueError(
+                    "Total measurement duration could not be determined from loops"
+                )
 
-        return self.metadata['num_frames'] / (total_duration/1000.0)
+        return self.metadata["num_frames"] / (total_duration / 1000.0)
 
     def _get_metadata_property(self, key, default=None):
         if self.metadata is None:
@@ -152,12 +164,22 @@ class ND2Reader(FramesSequenceND):
         """Setup the xyctz axes, iterate over t axis by default
 
         """
-        self._init_axis_if_exists('x', self._get_metadata_property("width", default=0))
-        self._init_axis_if_exists('y', self._get_metadata_property("height", default=0))
-        self._init_axis_if_exists('c', len(self._get_metadata_property("channels", default=[])), min_size=2)
-        self._init_axis_if_exists('t', len(self._get_metadata_property("frames", default=[])))
-        self._init_axis_if_exists('z', len(self._get_metadata_property("z_levels", default=[])), min_size=2)
-        self._init_axis_if_exists('v', len(self._get_metadata_property("fields_of_view", default=[])), min_size=2)
+        self._init_axis_if_exists("x", self._get_metadata_property("width", default=0))
+        self._init_axis_if_exists("y", self._get_metadata_property("height", default=0))
+        self._init_axis_if_exists(
+            "c", len(self._get_metadata_property("channels", default=[])), min_size=2
+        )
+        self._init_axis_if_exists(
+            "t", len(self._get_metadata_property("frames", default=[]))
+        )
+        self._init_axis_if_exists(
+            "z", len(self._get_metadata_property("z_levels", default=[])), min_size=2
+        )
+        self._init_axis_if_exists(
+            "v",
+            len(self._get_metadata_property("fields_of_view", default=[])),
+            min_size=2,
+        )
 
         if len(self.sizes) == 0:
             raise EmptyFileError("No axes were found for this .nd2 file.")
@@ -165,7 +187,7 @@ class ND2Reader(FramesSequenceND):
         # provide the default
         self.iter_axes = self._guess_default_iter_axis()
 
-        self._register_get_frame(self.get_frame_2D, 'yx')
+        self._register_get_frame(self.get_frame_2D, "yx")
 
     def _init_axis_if_exists(self, axis, size, min_size=1):
         if size >= min_size:
@@ -177,7 +199,7 @@ class ND2Reader(FramesSequenceND):
         Returns:
             the axis to iterate over
         """
-        priority = ['t', 'z', 'c', 'v']
+        priority = ["t", "z", "c", "v"]
         found_axes = []
         for axis in priority:
             try:
@@ -202,6 +224,9 @@ class ND2Reader(FramesSequenceND):
         if self._timesteps is not None and len(self._timesteps) > 0:
             return self._timesteps
 
-        self._timesteps = np.array(list(self._parser._raw_metadata.acquisition_times), dtype=np.float) * 1000.0
+        self._timesteps = (
+            np.array(list(self._parser._raw_metadata.acquisition_times), dtype=np.float)
+            * 1000.0
+        )
 
         return self._timesteps

tests/test_reader.py

@@ -16,18 +16,27 @@ class TestReader(unittest.TestCase):
     def test_extension(self):
         self.assertTrue('nd2' in ND2Reader.class_exts())
 
+    def cmp_two_readers(self, r1, r2):
+        attributes = r1.data['image_attributes']['SLxImageAttributes']
+        self.assertEqual(r2.metadata['width'], attributes['uiWidth'])
+        self.assertEqual(r2.metadata['height'], attributes['uiHeight'])
+
+        self.assertEqual(r2.metadata['width'], r2.sizes['x'])
+        self.assertEqual(r2.metadata['height'], r2.sizes['y'])
+
+        self.assertEqual(r2.pixel_type, np.float64)
+        self.assertEqual(r2.iter_axes, ['t'])
+
     def test_init_and_init_axes(self):
         with ArtificialND2('test_data/test_nd2_reader.nd2') as artificial:
             with ND2Reader('test_data/test_nd2_reader.nd2') as reader:
-                attributes = artificial.data['image_attributes']['SLxImageAttributes']
-                self.assertEqual(reader.metadata['width'], attributes['uiWidth'])
-                self.assertEqual(reader.metadata['height'], attributes['uiHeight'])
+                self.cmp_two_readers(artificial, reader)
 
-                self.assertEqual(reader.metadata['width'], reader.sizes['x'])
-                self.assertEqual(reader.metadata['height'], reader.sizes['y'])
-
-                self.assertEqual(reader.pixel_type, np.float64)
-                self.assertEqual(reader.iter_axes, ['t'])
+    def test_init_from_handler(self):
+        with ArtificialND2('test_data/test_nd2_reader.nd2') as artificial:
+            with open('test_data/test_nd2_reader.nd2', "rb") as FH:
+                with ND2Reader(FH) as reader:
+                    self.cmp_two_readers(artificial, reader)
 
     def test_init_empty_file(self):
         with ArtificialND2('test_data/empty.nd2', skip_blocks=['label_map_marker']):