Merge pull request #3 from danila-schelkov/sc2-support

refactor: inheritance removed, added support for the fifth version

Author: danila-schelkov

examples/compressor/main.py

@@ -3,17 +3,17 @@
 from sc_compression.signatures import Signatures
 from sc_compression import compress
 
-if not os.path.exists('in'):
-	os.mkdir('in')
-	
-if not os.path.exists('out'):
-	os.mkdir('out')
+if not os.path.exists("in"):
+    os.mkdir("in")
 
+if not os.path.exists("out"):
+    os.mkdir("out")
 
-for filename in os.listdir('in'):
-	with open('in/' + filename, 'rb') as f:
-		file_data = f.read()
-		f.close()
-	with open('out/' + filename, 'wb') as f:
-		f.write(compress(file_data, Signatures.SC, 3))
-		f.close()
+
+for filename in os.listdir("in"):
+    with open("in/" + filename, "rb") as f:
+        file_data = f.read()
+        f.close()
+    with open("out/" + filename, "wb") as f:
+        f.write(compress(file_data, Signatures.SC, 3))
+        f.close()

examples/decompressor/main.py

@@ -2,17 +2,17 @@
 
 from sc_compression import decompress
 
-if not os.path.exists('in'):
-	os.mkdir('in')
-	
-if not os.path.exists('out'):
-	os.mkdir('out')
+if not os.path.exists("in"):
+    os.mkdir("in")
 
+if not os.path.exists("out"):
+    os.mkdir("out")
 
-for filename in os.listdir('in'):
-	with open('in/' + filename, 'rb') as f:
-		file_data = f.read()
-		f.close()
-	with open('out/' + filename, 'wb') as f:
-		f.write(decompress(file_data)[0])
-		f.close()
+
+for filename in os.listdir("in"):
+    with open("in/" + filename, "rb") as f:
+        file_data = f.read()
+        f.close()
+    with open("out/" + filename, "wb") as f:
+        f.write(decompress(file_data)[0])
+        f.close()

sc_compression/__init__.py

@@ -2,23 +2,15 @@
 from .decompressor import Decompressor
 from .signatures import Signatures
 
-
-__all__ = [
-    'Decompressor',
-    'Compressor',
-    'Signatures',
-
-    'decompress',
-    'compress'
-]
+__all__ = ["Decompressor", "Compressor", "Signatures", "decompress", "compress"]
 
 
 def compress(buffer: bytes, signature: Signatures, file_version: int = None) -> bytes:
     return Compressor().compress(buffer, signature, file_version)
 
 
-def decompress(buffer: bytes) -> (bytes, int):
+def decompress(buffer: bytes) -> tuple[bytes, Signatures, int]:
     _decompressor = Decompressor()
     decompressed = _decompressor.decompress(buffer)
 
-    return decompressed, _decompressor.signature
+    return decompressed, _decompressor.signature, _decompressor.file_version

sc_compression/__main__.py

@@ -7,31 +7,53 @@
 
 
 def main() -> int:
-    argument_parser = argparse.ArgumentParser(description='A module for compression like in Supercell games.')
-    argument_parser.add_argument('-d', '--decompress', help='Decompress all files.', action='store_true')
-    argument_parser.add_argument('-c', '--compress', help='Compress all files.', action='store_true')
+    argument_parser = argparse.ArgumentParser(
+        description="A module for compression like in Supercell games."
+    )
+    argument_parser.add_argument(
+        "-d", "--decompress", help="Decompress all files", action="store_true"
+    )
     argument_parser.add_argument(
-        '-s', '--signature',
+        "-c", "--compress", help="Compress all files", action="store_true"
+    )
+    argument_parser.add_argument(
+        "-s",
+        "--signature",
         choices=[signature.name for signature in Signatures],
+        help="Required signature name",
+        type=str,
+        default=Signatures.SC.name,
+    )
+    argument_parser.add_argument(
+        "-f",
+        "--files",
+        help="Files to be processed",
+        type=str,
+        nargs="+",
+        required=True,
+    )
+    argument_parser.add_argument(
+        "-S",
+        "--suffix",
+        help="Suffix for processed files",
         type=str,
-        default='SC'
+        default=None,
     )
-    argument_parser.add_argument('-f', '--files', help='Files to be processed.', type=str, nargs='+', required=True)
 
     args = argument_parser.parse_args()
     if args.decompress:
         process_files(
             args.files,
             lambda data: decompress(data)[0],
-            'decomp'
+            args.suffix or "decompressed",
         )
     elif args.compress:
         signature = Signatures[args.signature]
 
         process_files(
             args.files,
             lambda data: compress(data, signature, 1),
-            'comp'
+            args.suffix or "compressed",
         )
     else:
         argument_parser.print_help()
@@ -58,25 +80,25 @@ def get_extension(file_path: str) -> str:
     return os.path.splitext(file_path)[1][1:]
 
 
-def join_filename(*parts, divider: str = '.') -> str:
+def join_filename(*parts, divider: str = ".") -> str:
     return divider.join(part for part in parts if type(part) is str and len(part) > 0)
 
 
-def process_files(files: str, action: Callable[[bytes], bytes], suffix: str):
+def process_files(files: list[str], action: Callable[[bytes], bytes], suffix: str):
     for filename in files:
         path = get_absolute_file_path(filename)
         directory = get_directory(path)
         base_name = get_base_name(path)
         extension = get_extension(path)
 
         if os.path.isfile(path):
-            with open(path, 'rb') as file:
+            with open(path, "rb") as file:
                 data = action(file.read())
 
             processed_filename = join_filename(base_name, suffix, extension)
-            with open(os.path.join(directory, processed_filename), 'wb') as file:
+            with open(os.path.join(directory, processed_filename), "wb") as file:
                 file.write(data)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     exit(main())

sc_compression/compressor.py

@@ -10,7 +10,7 @@
     from platform import system as get_system_name
 
     lzham = None
-    if get_system_name() == 'Windows':
+    if get_system_name() == "Windows":
         from sc_compression.support.lzham import LZHAM
 
         lzham = LZHAM
@@ -21,76 +21,90 @@
     zstandard = None
 
 
-class Compressor(Writer):
-    lzham_filters = {
-        'dict_size_log2': 18
-    }
+class Compressor:
+    lzham_filters = {"dict_size_log2": 18}
     lzma_filters = [
         {
             "id": lzma.FILTER_LZMA1,
             "dict_size": 256 * 1024,
             "lc": 3,
             "lp": 0,
             "pb": 2,
-            "mode": lzma.MODE_NORMAL
+            "mode": lzma.MODE_NORMAL,
         },
     ]
 
-    def __init__(self):
-        super().__init__('little')
-
-    def compress(self, data, signature: Signatures, file_version: int = None) -> bytes:
+    def compress(
+        self, data: bytes, signature: Signatures, file_version: int = None
+    ) -> bytes:
         uncompressed_size = len(data)
 
         if file_version is None:
             file_version = 3 if zstandard and signature != Signatures.SCLZ else 1
 
-        if signature == Signatures.ZSTD and not zstandard or \
-                signature == Signatures.SCLZ and not lzham:
+        if (
+            signature == Signatures.ZSTD
+            and not zstandard
+            or signature == Signatures.SCLZ
+            and not lzham
+        ):
             signature = Signatures.SC
 
-        super().__init__('little')
+        writer = Writer("little")
         if signature is Signatures.NONE:
             return data
-        elif signature in (Signatures.LZMA, Signatures.SIG) or (signature == Signatures.SC and file_version != 3):
-            compressed = lzma.compress(data, format=lzma.FORMAT_ALONE, filters=self.lzma_filters)
+        elif signature in (Signatures.LZMA, Signatures.SIG) or (
+            signature == Signatures.SC and file_version != 3
+        ):
+            compressed = lzma.compress(
+                data, format=lzma.FORMAT_ALONE, filters=self.lzma_filters
+            )
 
-            self.write(compressed[:5])
+            writer.write(compressed[:5])
 
-            self.writeInt32(uncompressed_size)
+            writer.write_int32(uncompressed_size)
 
-            self.write(compressed[13:])
+            writer.write(compressed[13:])
 
-            compressed = self.buffer
+            compressed = writer.buffer
         elif signature == Signatures.SCLZ and lzham:
             compressed = lzham.compress(data, filters=self.lzham_filters)
 
-            self.write(b'SCLZ')
-            self.writeUByte(self.lzham_filters['dict_size_log2'])
-            self.writeInt32(uncompressed_size)
-            self.write(compressed)
+            writer.write(b"SCLZ")
+            writer.write_u_int8(self.lzham_filters["dict_size_log2"])
+            writer.write_int32(uncompressed_size)
+            writer.write(compressed)
 
-            compressed = self.buffer
-        elif signature in (Signatures.SC, Signatures.ZSTD) and file_version == 3:
+            compressed = writer.buffer
+        elif signature in (Signatures.SC, Signatures.ZSTD) and zstandard is not None:
             compressor = zstandard.ZstdCompressor()
             compressed = compressor.compress(data)
         else:
-            raise TypeError('Unknown Signature!')
+            raise TypeError("Unknown Signature!")
+
+        compressed = self._write_header(compressed, data, file_version, signature)
+
+        return compressed
 
-        super().__init__('big')
+    @staticmethod
+    def _write_header(
+        compressed: bytes, data: bytes, file_version: int, signature: Signatures
+    ) -> bytes:
         if signature in (Signatures.SC, Signatures.SCLZ):
             data_hash = md5(data).digest()
 
-            self.write(b'SC')
-            self.writeInt32(file_version)
+            writer = Writer("big")
+            writer.write(b"SC")
+            writer.write_int32(file_version)
             if file_version == 4:
-                self.writeInt32(1)
-            self.writeInt32(len(data_hash))
-            self.write(data_hash)
-            compressed = self.buffer + compressed
+                writer.write_int32(1)
+            writer.write_int32(len(data_hash))
+            writer.write(data_hash)
+            return writer.buffer + compressed
         elif signature == Signatures.SIG:
-            self.write(b'Sig:')
-            self.write(b'\x00' * 64)  # sha64
-            compressed = self.buffer + compressed
+            writer = Writer("big")
+            writer.write(b"Sig:")
+            writer.write(b"\x00" * 64)  # sha64
+            return writer.buffer + compressed
 
         return compressed