[parser] Raise a RuntimeError if compression libs are not installed

m1stadev · Aug 22, 2023 · 60a885c · 60a885c
1 parent 63eb3b8
commit 60a885c
Showing 1 changed file with 30 additions and 2 deletions.
diff --git a/pyimg4/_parser.py b/pyimg4/_parser.py
@@ -2,13 +2,25 @@
 from zlib import adler32
 
 import asn1
-import liblzfse
-import lzss
 from Crypto.Cipher import AES
 
 from ._types import *
 from .errors import *
 
+try:
+    import lzss
+
+    have_lzss = True
+except ImportError:
+    have_lzss = False
+
+try:
+    import liblzfse
+
+    have_lzfse = True
+except ImportError:
+    have_lzfse = False
+
 
 class _PyIMG4:
     def __init__(self, data: Optional[bytes] = None) -> None:
@@ -1247,12 +1259,20 @@ def compress(self, compression: Compression) -> None:
             )
 
         if compression == Compression.LZSS:
+            if not have_lzss:
+                raise RuntimeError('pylzss not installed, cannot use LZSS compression')
+
             self._data = self._create_complzss_header() + lzss.compress(self._data)
 
             if self.extra is not None:
                 self._data += self.extra
 
         elif compression == Compression.LZFSE:
+            if not have_lzfse:
+                raise RuntimeError(
+                    'pyliblzfse not installed, cannot use LZFSE compression'
+                )
+
             payload_size = len(self._data)
             self._data = liblzfse.compress(self._data)
             # Cannot set LZFSE payload size until after compression
@@ -1275,10 +1295,18 @@ def decompress(self) -> None:
             raise CompressionError('Cannot decompress encrypted payload.')
 
         elif self.compression == Compression.LZSS:
+            if not have_lzss:
+                raise RuntimeError('pylzss not installed, cannot use LZSS compression')
+
             self._parse_complzss_header()
             self._data = lzss.decompress(self._data)
 
         elif self.compression == Compression.LZFSE:
+            if not have_lzfse:
+                raise RuntimeError(
+                    'pyliblzfse not installed, cannot use LZFSE compression'
+                )
+
             self._lzfse_payload_size = None
             self._data = liblzfse.decompress(self._data)