added user data tests; fixed writeUserData() with non-file streams

idelib/userdata.py

@@ -4,6 +4,7 @@
 the display of the `Dataset`.
 """
 
+import errno
 import os.path
 import logging
 from typing import Any, Dict, Optional, Tuple, Union
@@ -124,8 +125,9 @@ def writeUserData(dataset: Dataset,
         :param dataset: The `Dataset` from which to read the user data.
         :param userdata: A dictionary of user data, or `None` to remove
             existing user data. Note that the file will not get smaller if
-            the user data is removed (or the new user data is smaller);
-            it is just overwritten with null data (an EBML `Void` element).
+            the user data is removed or the new set of user data is smaller
+            than existing user data); it is just overwritten with null data
+            (an EBML `Void` element).
         :param refresh: If `True`, ignore any cached values and find the
             position in the file to which to write.
     """
@@ -145,6 +147,7 @@ def writeUserData(dataset: Dataset,
                                             lengthSize=8)
         else:
             # No new userdata, just write 'Void' over any existing userdata
+            # (or do nothing if there is no existing userdata)
             dataset._userdata = userdata
             if not hasdata:
                 return
@@ -154,18 +157,27 @@ def writeUserData(dataset: Dataset,
 
         userblob = dataBin + voidBin + offsetBin
 
-        if '+' not in fs.mode and 'w' not in fs.mode:
+        try:
+            writable = fs.writable()
+        except AttributeError:
+            # In case file-like stream doesn't implement `writable()`
+            # (e.g., older `ebmlite.threaded_file.ThreadAwareFile`)
+            mode = getattr(fs, 'mode', '')
+            writable = '+' in mode or 'w' in mode
+
+        if not writable:
             # File/stream is read-only; attempt to create a new file stream.
             if not getattr(fs, 'name', None):
-                logger.debug(f'(userdata) Dataset stream read only (mode {fs.mode!r}) '
-                             'and has no name, not writing user data')
-                return
+                raise IOError(errno.EACCES,
+                              f'Could not write user data; '
+                              f'Dataset stream not writable and has no filename')
 
             with open(fs.name, 'br+') as newfs:
                 logger.debug(f'(userdata) Dataset stream read only (mode {fs.mode!r}), '
                              'using new stream')
                 newfs.seek(offset, os.SEEK_SET)
                 newfs.write(userblob)
+
         else:
             fs.seek(offset, os.SEEK_SET)
             fs.write(userblob)

testing/file_streams.py

@@ -2,7 +2,9 @@
 
 FILES = [('./testing/SSX70065.IDE', 'rb'),
          ('./testing/SSX66115.IDE', 'rb'),
-         ('./test.ide', 'rb')]
+         ('./test.ide', 'rb'),
+         ('./testing/SSX_Data.IDE', 'rb'),
+         ('./testing/with_userdata.IDE', 'rb')]
 FILE_DICT = {}
 
 for fName, mode in FILES:

testing/test_userdata.py

@@ -0,0 +1,133 @@
+"""
+Test reading/writing user data to/from IDE files and streams (files and
+file-like).
+"""
+
+import pytest  # type: ignore
+
+from io import BytesIO
+import os.path
+import shutil
+
+from idelib import importer
+from idelib import userdata
+
+from testing import file_streams
+
+
+# ==============================================================================
+#
+# ==============================================================================
+
+USERDATA = {
+    'TimebaseOffset': 12345,
+    'WindowLayout': bytearray(b'bogus binary blob'),
+    'TimeBaseUTC': [1712769739]
+}
+
+SMALLER_USERDATA = {
+    'TimebaseOffset': 54321,
+}
+
+LARGER_USERDATA = {
+    'TimebaseOffset': 56789,
+    'WindowLayout': bytearray(b'bogus binary blob'),
+    'AnnotationList': {
+        'Annotation': [{'AnnotationID': 42, 'AnnotationStartTime': 101},],
+    },
+    'TimeBaseUTC': [35096400]
+}
+
+FILE_WITHOUT_USERDATA = './testing/SSX_Data.IDE'
+FILE_WITH_USERDATA = './testing/with_userdata.IDE'
+
+
+# ==============================================================================
+#
+# ==============================================================================
+
+def test_read_userdata():
+    """ Test reading user data.
+    """
+    doc = importer.openFile(file_streams.makeStreamLike(FILE_WITH_USERDATA))
+    data = userdata.readUserData(doc)
+    assert data == USERDATA
+
+
+def test_read_userdata_no_userdata():
+    """ Test reading user data from a file without user data.
+    """
+    doc = importer.openFile(file_streams.makeStreamLike(FILE_WITHOUT_USERDATA))
+    data = userdata.readUserData(doc)
+    assert data is None
+
+
+def test_write_userdata(tmp_path):
+    """ Test writing (and re-reading) user data to a file without existing
+        user data.
+    """
+    sourceFile = FILE_WITHOUT_USERDATA
+    filename = tmp_path / os.path.basename(sourceFile)
+
+    shutil.copyfile(sourceFile, filename)
+
+    with importer.importFile(filename) as doc:
+        userdata.writeUserData(doc, USERDATA)
+
+    with importer.importFile(filename) as doc:
+        data = userdata.readUserData(doc)
+        assert data == USERDATA
+
+
+def test_write_userdata_BytesIO():
+    """ Test writing (and re-reading) user data from a non-file stream
+        without existing user data.
+    """
+    sourceFile = FILE_WITHOUT_USERDATA
+
+    with open(sourceFile, 'rb') as f:
+        stream = BytesIO(f.read())
+
+    with importer.openFile(stream) as doc:
+        userdata.writeUserData(doc, USERDATA)
+
+        data = userdata.readUserData(doc)
+        assert data == USERDATA
+
+
+def test_larger_userdata(tmp_path):
+    """ Test overwriting an existing set of user data with a larger one.
+    """
+    sourceFile = FILE_WITH_USERDATA
+    filename = tmp_path / os.path.basename(sourceFile)
+    shutil.copyfile(sourceFile, filename)
+
+    originalSize = os.path.getsize(filename)
+
+    with importer.importFile(filename) as doc:
+        userdata.writeUserData(doc, LARGER_USERDATA)
+
+    with importer.importFile(filename) as doc:
+        data = userdata.readUserData(doc)
+        assert data == LARGER_USERDATA
+
+    assert originalSize < os.path.getsize(filename)
+
+
+def test_smaller_userdata(tmp_path):
+    """ Test overwriting an existing set of user data with a smaller one.
+    """
+    sourceFile = FILE_WITH_USERDATA
+    filename = tmp_path / os.path.basename(sourceFile)
+    shutil.copyfile(sourceFile, filename)
+
+    originalSize = os.path.getsize(filename)
+
+    with importer.importFile(filename) as doc:
+        userdata.writeUserData(doc, SMALLER_USERDATA)
+
+    with importer.importFile(filename) as doc:
+        data = userdata.readUserData(doc)
+        assert data == SMALLER_USERDATA
+
+    assert originalSize == os.path.getsize(filename)