diff --git a/src/RawFileReader.c b/src/RawFileReader.c index 473c72c..0b3084e 100644 --- a/src/RawFileReader.c +++ b/src/RawFileReader.c @@ -136,12 +136,18 @@ static PyObject *RawFileReader_read(RawFileReader *self, PyObject *args) { // resize the array to match the number of clusters read PyArray_Resize((PyArrayObject *)frames, &new_shape, 1, NPY_ANYORDER); + if(digital_frames){ + PyArray_Resize((PyArrayObject *)digital_frames, &new_shape, 1, NPY_ANYORDER); + } + // if we also read header we need to reshape the header if (self->read_header) { new_shape.len = 1; PyArray_Resize((PyArrayObject *)header, &new_shape, 1, NPY_ANYORDER); } + + } // Build up a tuple with the return values diff --git a/tests/test_RawFileReader.py b/tests/test_RawFileReader.py index 807da64..015676d 100644 --- a/tests/test_RawFileReader.py +++ b/tests/test_RawFileReader.py @@ -66,4 +66,13 @@ def test_references_on_m04_ad_with_header(data_path): assert sys.getrefcount(analog) == 2 #Over counts by one due to call by reference assert analog.shape == (3,400,400) assert sys.getrefcount(digital) == 2 - assert sys.getrefcount(header) == 2 \ No newline at end of file + assert sys.getrefcount(header) == 2 + +def test_read_too_many_m04_frames(data_path): + fname= data_path/'Moench04_digital_d0_f0_0.raw' + r = RawFileReader(fname, m04ad, header = True) + analog, digital, header = r.read(150) + assert analog.shape == (100,400,400) + assert digital.shape == (100, 400, 400) + assert header.size == 100 + assert sys.getrefcount(digital) == 2 \ No newline at end of file