Adds the StreamingBinWriter class that allows bin data to be streamed…

mffpy/bin_writer.py

@@ -14,7 +14,7 @@
 """
 from os import SEEK_SET
 from io import BytesIO, FileIO
-from typing import List, Union
+from typing import List, Union, IO
 from os.path import join
 
 import numpy as np
@@ -44,7 +44,7 @@ def __init__(self, sampling_rate: int, data_type: str = 'EEG'):
         self.data_type = data_type
         self.sampling_rate = sampling_rate
         self.header: Union[HeaderBlock, None] = None
-        self.stream = BytesIO()
+        self.stream: Union[IO[bytes], FileIO] = BytesIO()
         self.epochs: List[Epoch] = []
 
     @property
@@ -134,6 +134,7 @@ def write(self, filename: str, *args, **kwargs):
         # *args, **kwargs are ignored
         self.stream.seek(0, SEEK_SET)
         byts = self.stream.read()
+        assert isinstance(byts, bytes) 
         with open(filename, 'wb') as fo:
             num_written = fo.write(byts)
         assert num_written == len(byts), f"""
@@ -146,8 +147,25 @@ class StreamingBinWriter(BinWriter):
     """
 
     def __init__(self, sampling_rate: int, mffdir: str, data_type: str = 'EEG'):
+        """
+
+        **Parameters**
+
+        * **`sampling_rate`**: sampling rate of all channels.  Sampling rate
+        has to fit in a 3-byte integer.  See docs in `mffpy.header_block`.
+
+        * **`data_type`**: name of the type of signal.
+
+        * **`mffdir`**: directory of the mff recording to stream data to.
+
+        Note: Because we are streaming the recording to disk, the folder into which it
+        is to be saved must have been created prior to the initialization of this class.
+        """
+
         super().__init__(sampling_rate, data_type)
         self.stream = FileIO(join(mffdir, self.default_filename), mode='w')
 
     def write(self, filename: str, *args, **kwargs):
+        # Because the recording has been streamed to a file, all that is required 
+        # here is closing the stream
         self.stream.close()

mffpy/header_block.py

@@ -42,8 +42,9 @@
 
 import struct
 from os import SEEK_CUR
-from typing import IO
+from typing import IO, Union
 from collections import namedtuple
+from io import FileIO
 
 import numpy as np
 
@@ -141,7 +142,7 @@ def skip(n: int):
     )
 
 
-def write_header_block(fp: IO[bytes], hdr: HeaderBlock):
+def write_header_block(fp: Union[IO[bytes], FileIO], hdr: HeaderBlock):
     """write HeaderBlock `hdr` to file pointer `fp`"""
     fp.write(struct.pack('4i',
                          1, hdr.header_size, hdr.block_size, hdr.num_channels))

mffpy/tests/test_writer.py

@@ -112,8 +112,6 @@ def test_writer_exports_JSON():
     except BaseException:
         raise AssertionError(f"""Clean-up failed of '{filename}'.""")
 
-##START##
-
 def test_streaming_writer_receives_bad_init_data():
     """Test bin writer fails when initialized with non-int sampling rate"""
     dirname = 'testdir.mff'
@@ -124,45 +122,39 @@ def test_streaming_writer_receives_bad_init_data():
     rmtree(dirname)
 
 def test_streaming_writer_writes():
-    dirname = 'testdir2.mff'
+    dirname = 'testdir3.mff'
     # create some data and add it to a binary writer
     device = 'HydroCel GSN 256 1.0'
     num_samples = 10
     num_channels = 256
     sampling_rate = 128
     # create an mffpy.Writer and add a file info, and the binary file
-    W = Writer(dirname)
-    W.create_directory()
-    b = StreamingBinWriter(sampling_rate=sampling_rate, data_type='EEG', mffdir=dirname)
+    writer = Writer(dirname)
+    writer.create_directory()
+    bin_writer = StreamingBinWriter(sampling_rate=sampling_rate, data_type='EEG', mffdir=dirname)
     data = np.random.randn(num_channels, num_samples).astype(np.float32)
-    b.add_block(data)
+    bin_writer.add_block(data)
     startdatetime = datetime.strptime(
         '1984-02-18T14:00:10.000000+0100', XML._time_format)
-    W.addxml('fileInfo', recordTime=startdatetime)
-    W.add_coordinates_and_sensor_layout(device)
-    W.addbin(b)
-    W.write()
+    writer.addxml('fileInfo', recordTime=startdatetime)
+    writer.add_coordinates_and_sensor_layout(device)
+    writer.addbin(bin_writer)
+    writer.write()
     # read it again; compare the result
-    R = Reader(dirname)
-    assert R.startdatetime == startdatetime
+    reader = Reader(dirname)
+    assert reader.startdatetime == startdatetime
     # Read binary data and compare
-    read_data = R.get_physical_samples_from_epoch(R.epochs[0])
+    read_data = reader.get_physical_samples_from_epoch(reader.epochs[0])
     assert 'EEG' in read_data
     read_data, t0 = read_data['EEG']
     assert t0 == 0.0
     assert read_data == pytest.approx(data)
-    layout = R.directory.filepointer('sensorLayout')
+    layout = reader.directory.filepointer('sensorLayout')
     layout = XML.from_file(layout)
     assert layout.name == device
     # cleanup
     try:
-        remove(join(dirname, 'info.xml'))
-        remove(join(dirname, 'info1.xml'))
-        remove(join(dirname, 'epochs.xml'))
-        remove(join(dirname, 'signal1.bin'))
-        remove(join(dirname, 'coordinates.xml'))
-        remove(join(dirname, 'sensorLayout.xml'))
-        rmdir(dirname)
+        rmtree(dirname)
     except BaseException:
         raise AssertionError(f"""
         Clean-up failed of '{dirname}'.  Were additional files written?""")

mffpy/writer.py

@@ -39,6 +39,7 @@ def __init__(self, filename: str):
         self.file_created = False
 
     def create_directory(self):
+        """Creates the directory for the recording."""
         if not self.file_created:
             makedirs(self.mffdir, exist_ok=False)
             self.file_created = True