Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit f4b1ba1

Browse files
addisonElliottaddisonElliott
authored
Support reading from BytesIO (#125)
Add logic to correctly handle reading from BytesIO objects. Previously, parsing the data used np.fromfile which does not support BytesIO objects. This is fixed by using np.frombuffer and np.fromstring to read in BytesIO objects. np.fromfile does not work with BytesIO objects and therefore caused Additionally, tests for reading/writing from BytesIO objects were enhanced.
1 parent 4f97d33 commit f4b1ba1

3 files changed

Lines changed: 165 additions & 41 deletions

File tree

‎nrrd/reader.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import bz2
2+
import io
23
import os
34
import re
45
import shlex
@@ -401,9 +402,16 @@ def read_data(header, fh=None, filename=None, index_order='F'):
401402

402403
# If a compression encoding is used, then byte skip AFTER decompressing
403404
if header['encoding'] == 'raw':
404-
data = np.fromfile(fh, dtype)
405+
if isinstance(fh, io.BytesIO):
406+
raw_data = bytearray(fh.read(total_data_points * dtype.itemsize))
407+
data = np.frombuffer(raw_data, dtype)
408+
else:
409+
data = np.fromfile(fh, dtype)
405410
elif header['encoding'] in ['ASCII', 'ascii', 'text', 'txt']:
406-
data = np.fromfile(fh, dtype, sep=' ')
411+
if isinstance(fh, io.BytesIO):
412+
data = np.fromstring(fh.read(), dtype, sep=' ')
413+
else:
414+
data = np.fromfile(fh, dtype, sep=' ')
407415
else:
408416
# Handle compressed data now
409417
# Construct the decompression object based on encoding

‎nrrd/tests/test_reading.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import io
12
import unittest
23
from typing import ClassVar
34

@@ -465,6 +466,41 @@ def test_read_quoted_string_header_no_quotes(self):
465466
self.assertEqual(['mm', 'cm', 'in'], header['space units'])
466467
self.assertEqual(['X', 'Y', 'f(log(X,10),Y)'], header['labels'])
467468

469+
def test_read_memory(self):
470+
def test(filename: str):
471+
with open(filename, 'rb') as fh:
472+
# Read into BytesIO and test that
473+
x = fh.read()
474+
memory_file = io.BytesIO(x)
475+
memory_header = nrrd.read_header(memory_file)
476+
memory_data = nrrd.read_data(memory_header, memory_file)
477+
478+
# Read normally via file handle
479+
fh.seek(0)
480+
expected_header = nrrd.read_header(fh)
481+
expected_data = nrrd.read_data(expected_header, fh)
482+
483+
np.testing.assert_equal(expected_header, memory_header)
484+
np.testing.assert_equal(expected_data, memory_data)
485+
486+
# Test that the data read is able to be edited
487+
self.assertTrue(memory_data.flags['WRITEABLE'])
488+
489+
paths = [
490+
RAW_NRRD_FILE_PATH,
491+
GZ_NRRD_FILE_PATH,
492+
GZ_BYTESKIP_NRRD_FILE_PATH,
493+
GZ_LINESKIP_NRRD_FILE_PATH,
494+
BZ2_NRRD_FILE_PATH,
495+
ASCII_1D_NRRD_FILE_PATH,
496+
ASCII_2D_NRRD_FILE_PATH,
497+
RAW_4D_NRRD_FILE_PATH,
498+
]
499+
500+
for filename in paths:
501+
with self.subTest(filename):
502+
test(filename)
503+
468504

469505
class TestReadingFunctionsFortran(Abstract.TestReadingFunctions):
470506
index_order = 'F'

‎nrrd/tests/test_writing.py‎

Lines changed: 119 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -21,44 +21,66 @@ def setUp(self):
2121
with open(RAW_DATA_FILE_PATH, 'rb') as fh:
2222
self.expected_data = fh.read()
2323

24-
def write_and_read_back(self, encoding=None, level=9):
25-
output_filename = os.path.join(self.temp_write_dir, f'testfile_{encoding}_{str(level)}.nrrd')
26-
headers = {}
27-
if encoding is not None:
28-
headers['encoding'] = encoding
29-
nrrd.write(output_filename, self.data_input, headers, compression_level=level,
30-
index_order=self.index_order)
24+
def test_write_default_header(self):
25+
output_filename = os.path.join(self.temp_write_dir, 'testfile_default.nrrd')
26+
nrrd.write(output_filename, self.data_input, {}, index_order=self.index_order)
3127

3228
# Read back the same file
3329
data, header = nrrd.read(output_filename, index_order=self.index_order)
3430
self.assertEqual(self.expected_data, data.tobytes(order=self.index_order))
35-
self.assertEqual(header.get('encoding'), encoding or 'gzip') # default is gzip is not specified
36-
37-
return output_filename
38-
39-
def test_write_default_header(self):
40-
self.write_and_read_back()
31+
self.assertEqual(header.get('encoding'), 'gzip') # default is gzip if not specified
4132

4233
def test_write_raw(self):
43-
self.write_and_read_back('raw')
34+
output_filename = os.path.join(self.temp_write_dir, 'testfile_raw.nrrd')
35+
nrrd.write(output_filename, self.data_input, {'encoding': 'raw'}, index_order=self.index_order)
36+
37+
# Read back the same file
38+
data, header = nrrd.read(output_filename, index_order=self.index_order)
39+
self.assertEqual(self.expected_data, data.tobytes(order=self.index_order))
40+
self.assertEqual(header.get('encoding'), 'raw')
4441

4542
def test_write_gz(self):
46-
self.write_and_read_back('gzip')
43+
output_filename = os.path.join(self.temp_write_dir, 'testfile_gzip.nrrd')
44+
nrrd.write(output_filename, self.data_input, {'encoding': 'gzip'}, index_order=self.index_order)
45+
46+
# Read back the same file
47+
data, header = nrrd.read(output_filename, index_order=self.index_order)
48+
self.assertEqual(self.expected_data, data.tobytes(order=self.index_order))
49+
self.assertEqual(header.get('encoding'), 'gzip')
4750

4851
def test_write_bzip2(self):
49-
self.write_and_read_back('bzip2')
52+
output_filename = os.path.join(self.temp_write_dir, 'testfile_bzip2.nrrd')
53+
nrrd.write(output_filename, self.data_input, {'encoding': 'bzip2'}, index_order=self.index_order)
54+
55+
# Read back the same file
56+
data, header = nrrd.read(output_filename, index_order=self.index_order)
57+
self.assertEqual(self.expected_data, data.tobytes(order=self.index_order))
58+
self.assertEqual(header.get('encoding'), 'bzip2')
5059

5160
def test_write_gz_level1(self):
52-
filename = self.write_and_read_back('gzip', level=1)
61+
output_filename = os.path.join(self.temp_write_dir, 'testfile_gzip_1.nrrd')
62+
nrrd.write(output_filename, self.data_input, {'encoding': 'gzip'}, compression_level=1,
63+
index_order=self.index_order)
5364

54-
self.assertLess(os.path.getsize(GZ_NRRD_FILE_PATH), os.path.getsize(filename))
65+
# Read back the same file
66+
data, header = nrrd.read(output_filename, index_order=self.index_order)
67+
self.assertEqual(self.expected_data, data.tobytes(order=self.index_order))
68+
self.assertEqual(header.get('encoding'), 'gzip')
69+
self.assertLess(os.path.getsize(GZ_NRRD_FILE_PATH), os.path.getsize(output_filename))
5570

5671
def test_write_bzip2_level1(self):
57-
_ = self.write_and_read_back('bzip2', level=1)
72+
output_filename = os.path.join(self.temp_write_dir, 'testfile_bzip2_1.nrrd')
73+
nrrd.write(output_filename, self.data_input, {'encoding': 'bzip2'}, compression_level=1,
74+
index_order=self.index_order)
75+
76+
# Read back the same file
77+
data, header = nrrd.read(output_filename, index_order=self.index_order)
78+
self.assertEqual(self.expected_data, data.tobytes(order=self.index_order))
79+
self.assertEqual(header.get('encoding'), 'bzip2')
5880

5981
# note: we don't currently assert reduction here, because with the binary ball test data,
6082
# the output size does not change at different bz2 levels.
61-
# self.assertLess(os.path.getsize(BZ2_NRRD_FILE_PATH), os.path.getsize(fn))
83+
# self.assertLess(os.path.getsize(BZ2_NRRD_FILE_PATH), os.path.getsize(output_filename))
6284

6385
def test_write_ascii_1d(self):
6486
output_filename = os.path.join(self.temp_write_dir, 'testfile_ascii_1d.nrrd')
@@ -357,25 +379,83 @@ def test_write_check_remove_datafile(self):
357379
data, header = nrrd.read(output_filename, index_order=self.index_order)
358380
self.assertFalse('data file' in header)
359381

360-
def test_write_memory(self):
361-
default_output_filename = os.path.join(self.temp_write_dir, 'testfile_default_filename.nrrd')
362-
nrrd.write(default_output_filename, self.data_input, {}, index_order=self.index_order)
363-
364-
memory_nrrd = io.BytesIO()
365-
366-
nrrd.write(memory_nrrd, self.data_input, {}, index_order=self.index_order)
367-
368-
memory_nrrd.seek(0)
369-
370-
data, header = nrrd.read(default_output_filename, index_order=self.index_order)
371-
memory_header = nrrd.read_header(memory_nrrd)
372-
memory_data = nrrd.read_data(header=memory_header, fh=memory_nrrd, filename=None,
373-
index_order=self.index_order)
374-
375-
self.assertEqual(self.expected_data, data.tobytes(order=self.index_order))
376-
self.assertEqual(self.expected_data, memory_data.tobytes(order=self.index_order))
377-
self.assertEqual(header.pop('sizes').all(), memory_header.pop('sizes').all())
378-
self.assertSequenceEqual(header, memory_header)
382+
def test_write_memory_default(self):
383+
kwargs = {
384+
'header': {},
385+
'index_order': self.index_order
386+
}
387+
388+
memory_nrrd_file = io.BytesIO()
389+
nrrd.write(memory_nrrd_file, self.data_input, **kwargs)
390+
memory_nrrd_file.seek(0)
391+
memory_nrrd = memory_nrrd_file.readlines()
392+
393+
expected_filename = os.path.join(self.temp_write_dir, 'testfile_expected.nrrd')
394+
nrrd.write(expected_filename, self.data_input, **kwargs)
395+
with open(expected_filename, 'rb') as fh:
396+
expected_nrrd = fh.readlines()
397+
398+
self.assertEqual(expected_nrrd, memory_nrrd)
399+
400+
def test_write_memory_raw(self):
401+
kwargs = {
402+
'header': {
403+
'encoding': 'raw'
404+
},
405+
'index_order': self.index_order
406+
}
407+
408+
memory_nrrd_file = io.BytesIO()
409+
nrrd.write(memory_nrrd_file, self.data_input, **kwargs)
410+
memory_nrrd_file.seek(0)
411+
memory_nrrd = memory_nrrd_file.readlines()
412+
413+
expected_filename = os.path.join(self.temp_write_dir, 'testfile_expected.nrrd')
414+
nrrd.write(expected_filename, self.data_input, **kwargs)
415+
with open(expected_filename, 'rb') as fh:
416+
expected_nrrd = fh.readlines()
417+
418+
self.assertEqual(expected_nrrd, memory_nrrd)
419+
420+
def test_write_memory_gzip(self):
421+
kwargs = {
422+
'header': {
423+
'encoding': 'gzip'
424+
},
425+
'index_order': self.index_order
426+
}
427+
428+
memory_nrrd_file = io.BytesIO()
429+
nrrd.write(memory_nrrd_file, self.data_input, **kwargs)
430+
memory_nrrd_file.seek(0)
431+
memory_nrrd = memory_nrrd_file.readlines()
432+
433+
expected_filename = os.path.join(self.temp_write_dir, 'testfile_expected.nrrd')
434+
nrrd.write(expected_filename, self.data_input, **kwargs)
435+
with open(expected_filename, 'rb') as fh:
436+
expected_nrrd = fh.readlines()
437+
438+
self.assertEqual(expected_nrrd, memory_nrrd)
439+
440+
def test_write_memory_bzip2(self):
441+
kwargs = {
442+
'header': {
443+
'encoding': 'bzip2'
444+
},
445+
'index_order': self.index_order
446+
}
447+
448+
memory_nrrd_file = io.BytesIO()
449+
nrrd.write(memory_nrrd_file, self.data_input, **kwargs)
450+
memory_nrrd_file.seek(0)
451+
memory_nrrd = memory_nrrd_file.readlines()
452+
453+
expected_filename = os.path.join(self.temp_write_dir, 'testfile_expected.nrrd')
454+
nrrd.write(expected_filename, self.data_input, **kwargs)
455+
with open(expected_filename, 'rb') as fh:
456+
expected_nrrd = fh.readlines()
457+
458+
self.assertEqual(expected_nrrd, memory_nrrd)
379459

380460
def test_write_memory_file_handle(self):
381461
default_output_filename = os.path.join(self.temp_write_dir, 'testfile_default_filename.nrrd')

0 commit comments

Comments
 (0)