Refactored _write function, added tests for bytes written

dissect/cstruct/cstruct.py

@@ -11,6 +11,7 @@
 from dissect.cstruct.expression import Expression
 from dissect.cstruct.parser import CStyleParser, TokenParser
 from dissect.cstruct.types import (
+    LEB128,
     Array,
     ArrayMetaType,
     BaseType,
@@ -29,7 +30,6 @@
     Wchar,
     WcharArray,
 )
-from dissect.cstruct.types.leb128 import LEB128
 
 
 class cstruct:

dissect/cstruct/types/leb128.py

@@ -47,7 +47,7 @@ def _read_0(cls, stream: BinaryIO, context: dict[str, Any] = None) -> LEB128:
         return result
 
     @classmethod
-    def _write(cls, stream: BinaryIO, data: int) -> LEB128:
+    def _write(cls, stream: BinaryIO, data: int) -> int:
         # only write negative numbers when in signed mode
         if data < 0 and not cls.signed:
             raise ValueError("Attempt to encode a negative integer using unsigned LEB128 encoding")
@@ -64,8 +64,10 @@ def _write(cls, stream: BinaryIO, data: int) -> LEB128:
                 not cls.signed and data == 0
             ):
                 result.append(byte)
-                stream.write(result)
-                return len(result)
+                break
 
             # Set high-order bit of byte
             result.append(0x80 | byte)
+
+        stream.write(result)
+        return len(result)

tests/test_types_leb128.py

@@ -1,3 +1,5 @@
+import io
+
 import pytest
 
 from dissect.cstruct.cstruct import cstruct
@@ -159,3 +161,17 @@ def test_leb128_write_negatives(cs: cstruct):
     with pytest.raises(ValueError, match="Attempt to encode a negative integer using unsigned LEB128 encoding"):
         cs.uleb128(-2).dumps()
     assert cs.ileb128(-2).dumps() == b"\x7e"
+
+
+def test_leb128_unsigned_write_amount_written(cs: cstruct):
+    out1 = io.BytesIO()
+    bytes_written1 = cs.uleb128(2).write(out1)
+    assert bytes_written1 == out1.tell()
+
+    out2 = io.BytesIO()
+    bytes_written2 = cs.uleb128(4747).write(out2)
+    assert bytes_written2 == out2.tell()
+
+    out3 = io.BytesIO()
+    bytes_written3 = cs.uleb128(13371337).write(out3)
+    assert bytes_written3 == out3.tell()