RFCT utils.fileish_open with a write possibility

Author: xmedeko

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()

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)

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)"""

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()