Merge pull request #16 from remopas/dev

fix class format again
remkop22 · May 19, 2022 · 9fa942e · 9fa942e
2 parents 5a88a88 + fc15bbd
commit 9fa942e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/src/binread/format.py b/src/binread/format.py
@@ -183,9 +183,14 @@ def __init__(self, cls: type, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._cls = cls
 
-    def extract(self, data: bytes, fields: Dict[str, Any]) -> Tuple[Any, int]:
-        value, bytes_read = super().extract(data, fields)
-        return self._cls(**value), bytes_read
+    def read(self, *args, **kwargs) -> Dict[str, Any]:
+        result = super().read(*args, **kwargs)
+
+        if isinstance(result, tuple):
+            result, bytes_read = result
+            return self._cls(**result), bytes_read  # type: ignore
+        else:
+            return self._cls(**result)
 
 
 def formatclass(*args, **kwargs):

diff --git a/tests/binread/test_format.py b/tests/binread/test_format.py
@@ -11,4 +11,5 @@ class FormatExample:
             field1 = binread.U8
             field2 = binread.U16
 
-        FormatExample.read(b"\x00\xFF\x00")
+        value = FormatExample.read(b"\x00\xFF\x00")
+        self.assertEqual(type(value), FormatExample)