Repository: psd-tools/psd-tools
Pull request #413 (2024-06-14 00:43:18)

Raise IO error instead of assertion
Changes:
- Raise IO error instead of assertion


Patch:
diff --git a/src/psd_tools/psd/adjustments.py b/src/psd_tools/psd/adjustments.py
index 7b7fd65..000c17a 100644
--- a/src/psd_tools/psd/adjustments.py
+++ b/src/psd_tools/psd/adjustments.py
@@ -204,7 +204,7 @@ class Curves(BaseElement):
         if version == 1:
             try:
                 extra = CurvesExtraMarker.read(fp, is_map=is_map)
-            except AssertionError:
+            except IOError:
                 logger.warning("Failed to read CurvesExtraMarker")
 
         return cls(is_map, version, count_map, data, extra)
diff --git a/src/psd_tools/psd/base.py b/src/psd_tools/psd/base.py
index d7c3cd5..fe37052 100644
--- a/src/psd_tools/psd/base.py
+++ b/src/psd_tools/psd/base.py
@@ -363,7 +363,7 @@ class ShortIntegerElement(IntegerElement):
     def read(cls, fp, **kwargs):
         try:
             return cls(read_fmt("H2x", fp)[0])
-        except AssertionError as e:
+        except IOError as e:
             logger.error(e)
         return cls(read_fmt("H", fp)[0])
 
@@ -383,7 +383,7 @@ class ByteElement(IntegerElement):
     def read(cls, fp, **kwargs):
         try:
             return cls(read_fmt("B3x", fp)[0])
-        except AssertionError as e:
+        except IOError as e:
             logger.error(e)
         return cls(read_fmt("B", fp)[0])
 
@@ -405,7 +405,7 @@ class BooleanElement(IntegerElement):
     def read(cls, fp, **kwargs):
         try:
             return cls(read_fmt("?3x", fp)[0])
-        except AssertionError as e:
+        except IOError as e:
             logger.error(e)
         return cls(read_fmt("?", fp)[0])
 
diff --git a/src/psd_tools/utils.py b/src/psd_tools/utils.py
index 47da807..068766c 100644
--- a/src/psd_tools/utils.py
+++ b/src/psd_tools/utils.py
@@ -34,11 +34,12 @@ def read_fmt(fmt, fp):
     fmt = str(">" + fmt)
     fmt_size = struct.calcsize(fmt)
     data = fp.read(fmt_size)
-    try:
-        assert len(data) == fmt_size, "read=%d, expected=%d" % (len(data), fmt_size)
-    except AssertionError:
+    if len(data) != fmt_size:
         fp.seek(-len(data), 1)
-        raise
+        raise IOError(
+            "Failed to read data section: read=%d, expected=%d. "
+            "Likely the file is corrupted." % (len(data), fmt_size)
+        )
     return struct.unpack(fmt, data)
 
 
@@ -49,7 +50,11 @@ def write_fmt(fp, fmt, *args):
     fmt = str(">" + fmt)
     fmt_size = struct.calcsize(fmt)
     written = write_bytes(fp, struct.pack(fmt, *args))
-    assert written == fmt_size, "written=%d, expected=%d" % (written, fmt_size)
+    if written != fmt_size:
+        raise IOError(
+            "Failed to write data section: written=%d, expected=%d."
+            % (written, fmt_size)
+        )
     return written
 
 
@@ -62,7 +67,10 @@ def write_bytes(fp, data):
     pos = fp.tell()
     fp.write(data)
     written = fp.tell() - pos
-    assert written == len(data), "written=%d, expected=%d" % (written, len(data))
+    if written != len(data):
+        raise IOError(
+            "Failed to write data: written=%d, expected=%d." % (written, len(data))
+        )
     return written
 
 
@@ -76,7 +84,11 @@ def read_length_block(fp, fmt="I", padding=1):
     """
     length = read_fmt(fmt, fp)[0]
     data = fp.read(length)
-    assert len(data) == length, (len(data), length)
+    if len(data) != length:
+        raise IOError(
+            "Failed to read data section: read=%d, expected=%d. "
+            "Likely the file is corrupted." % (len(data), length)
+        )
     read_padding(fp, length, padding)
     return data
 


Test patch:
diff --git a/tests/psd_tools/psd/test_color_mode_data.py b/tests/psd_tools/psd/test_color_mode_data.py
index 6673102..9862ddc 100644
--- a/tests/psd_tools/psd/test_color_mode_data.py
+++ b/tests/psd_tools/psd/test_color_mode_data.py
@@ -1,7 +1,6 @@
 from __future__ import absolute_import, unicode_literals
 
 import pytest
-
 from psd_tools.psd.color_mode_data import ColorModeData
 from psd_tools.utils import pack
 
@@ -93,7 +92,7 @@ def test_color_mode_data(fixture):
 
 
 def test_color_mode_data_exception():
-    with pytest.raises(AssertionError):
+    with pytest.raises(IOError):
         ColorModeData.frombytes(b"")
-    with pytest.raises(AssertionError):
+    with pytest.raises(IOError):
         ColorModeData.frombytes(str(10).encode("ascii"))
diff --git a/tests/psd_tools/psd/test_image_resources.py b/tests/psd_tools/psd/test_image_resources.py
index a16bab8..9ed4edf 100644
--- a/tests/psd_tools/psd/test_image_resources.py
+++ b/tests/psd_tools/psd/test_image_resources.py
@@ -16,7 +16,7 @@ def test_image_resources_from_to():
 
 
 def test_image_resources_exception():
-    with pytest.raises(AssertionError):
+    with pytest.raises(IOError):
         ImageResources.frombytes(b"\x00\x00\x00\x01")
 
 
@@ -40,7 +40,7 @@ def test_image_resource_from_to(fixture):
 
 
 def test_image_resource_exception():
-    with pytest.raises(AssertionError):
+    with pytest.raises(IOError):
         ImageResource.frombytes(b"\x00\x00\x00\x01")
 
 
diff --git a/tests/psd_tools/psd/test_psd.py b/tests/psd_tools/psd/test_psd.py
index 6d32955..f8f187e 100644
--- a/tests/psd_tools/psd/test_psd.py
+++ b/tests/psd_tools/psd/test_psd.py
@@ -4,7 +4,6 @@ import io
 import os
 
 import pytest
-
 from psd_tools.psd import PSD
 
 from ..utils import TEST_ROOT, all_files, check_write_read
@@ -54,7 +53,7 @@ def test_psd_write_read(filename):
 
 
 def test_psd_from_error():
-    with pytest.raises(AssertionError):
+    with pytest.raises(IOError):
         PSD.frombytes(b"\x00\x00\x00\x00")