diff --git a/fitparse/base.py b/fitparse/base.py index 4bcd494..a37d9e0 100644 --- a/fitparse/base.py +++ b/fitparse/base.py @@ -16,23 +16,12 @@ BASE_TYPES, BASE_TYPE_BYTE, DevField, add_dev_data_id, add_dev_field_description, get_dev_type ) -from fitparse.utils import calc_crc, FitParseError, FitEOFError, FitCRCError, FitHeaderError +from fitparse.utils import calc_crc, fileish_open, FitParseError, FitEOFError, FitCRCError, FitHeaderError + class FitFile(object): def __init__(self, fileish, check_crc=True, data_processor=None): - if hasattr(fileish, 'read'): - # BytesIO-like object - self._file = fileish - elif isinstance(fileish, str): - # Python2 - file path, file contents in the case of a TypeError - # Python3 - file path - try: - self._file = open(fileish, 'rb') - except TypeError: - self._file = io.BytesIO(fileish) - else: - # Python 3 - file contents - self._file = io.BytesIO(fileish) + self._file = fileish_open(fileish, 'rb') self.check_crc = check_crc self._processor = data_processor or FitFileDataProcessor() diff --git a/fitparse/utils.py b/fitparse/utils.py index 9ff2b73..caed865 100644 --- a/fitparse/utils.py +++ b/fitparse/utils.py @@ -1,5 +1,7 @@ import re +import io + class FitParseError(ValueError): pass @@ -47,3 +49,29 @@ def scrub_method_name(method_name, convert_units=False): replace_from, '%s' % replace_to, ) return METHOD_NAME_SCRUBBER.sub('_', method_name) + + +def fileish_open(fileish, mode): + """ + Convert file-ish object to BytesIO like object. + :param fileish: the file-ihs object (str, BytesIO, bytes, file contents) + :param str mode: mode for the open function. + :rtype: BytesIO + """ + if mode is not None and any(m in mode for m in ['+', 'w', 'a', 'x']): + attr = 'write' + else: + attr = 'read' + if hasattr(fileish, attr) and hasattr(fileish, 'seek'): + # BytesIO-like object + return fileish + elif isinstance(fileish, str): + # Python2 - file path, file contents in the case of a TypeError + # Python3 - file path + try: + return open(fileish, mode) + except TypeError: + return io.BytesIO(fileish) + else: + # Python 3 - file contents + return io.BytesIO(fileish) diff --git a/tests/test.py b/tests/test.py index 0149d25..8b34082 100755 --- a/tests/test.py +++ b/tests/test.py @@ -398,21 +398,6 @@ def test_int_long(self): with FitFile(testfile('event_timestamp.fit')) as f: assert f.messages[-1].fields[1].raw_value == 1739.486328125 - def test_fileish_types(self): - """Test the constructor does the right thing when given different types - (specifically, test files with 8 characters, followed by an uppercase.FIT - extension), which confused the fileish check on Python 2, see - https://github.com/dtcooper/python-fitparse/issues/29#issuecomment-312436350 - for details""" - with FitFile(testfile('nametest.FIT')): - pass - with open(testfile("nametest.FIT"), 'rb') as f: - FitFile(f) - with open(testfile("nametest.FIT"), 'rb') as f: - FitFile(f.read()) - with open(testfile("nametest.FIT"), 'rb') as f: - FitFile(io.BytesIO(f.read())) - def test_elemnt_bolt_developer_data_id_without_application_id(self): """Test that a file without application id set inside developer_data_id is parsed (as seen on ELEMNT BOLT with firmware version WB09-1507)""" diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..e0b6547 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python + +import io +import os +import sys +import tempfile + +from fitparse.utils import fileish_open + +if sys.version_info >= (2, 7): + import unittest +else: + import unittest2 as unittest + + +def testfile(filename): + return os.path.join(os.path.dirname(os.path.realpath(__file__)), 'files', filename) + + +class UtilsTestCase(unittest.TestCase): + + def test_fileish_open_read(self): + """Test the constructor does the right thing when given different types + (specifically, test files with 8 characters, followed by an uppercase.FIT + extension), which confused the fileish check on Python 2, see + https://github.com/dtcooper/python-fitparse/issues/29#issuecomment-312436350 + for details""" + + def test_fopen(fileish): + with fileish_open(fileish, 'rb') as f: + self.assertIsNotNone(f.read(1)) + f.seek(0, os.SEEK_SET) + + test_fopen(testfile('nametest.FIT')) + with open(testfile("nametest.FIT"), 'rb') as f: + test_fopen(f) + with open(testfile("nametest.FIT"), 'rb') as f: + test_fopen(f.read()) + with open(testfile("nametest.FIT"), 'rb') as f: + test_fopen(io.BytesIO(f.read())) + + def test_fileish_open_write(self): + + def test_fopen(fileish): + with fileish_open(fileish, 'wb') as f: + f.write(b'\x12') + f.seek(0, os.SEEK_SET) + + tmpfile = tempfile.NamedTemporaryFile(prefix='fitparse-test', suffix='.FIT', delete=False) + filename = tmpfile.name + tmpfile.close() + try: + test_fopen(filename) + with open(filename, 'wb') as f: + test_fopen(f) + test_fopen(io.BytesIO()) + finally: + # remove silently + try: + os.remove(filename) + except OSError: + pass + + +if __name__ == '__main__': + unittest.main()