Extract size attribute from ybegin line (#92)

src/sabctools.pyi

@@ -5,7 +5,7 @@ __version__: str
 openssl_linked: bool
 simd: str
 
-def yenc_decode(raw_data: bytearray) -> Tuple[str, int, int, Optional[int]]: ...
+def yenc_decode(raw_data: bytearray) -> Tuple[str, int, int, int, Optional[int]]: ...
 def yenc_encode(input_string: bytes) -> Tuple[bytes, int]: ...
 
 def unlocked_ssl_recv_into(ssl_socket: SSLSocket, buffer: memoryview) -> int: ...

src/yenc.cc

@@ -71,6 +71,7 @@ PyObject* yenc_decode(PyObject* self, PyObject* Py_bytesarray_obj) {
     uint32_t crc_yenc = 0;
     size_t yenc_data_length;
     size_t output_len;
+    unsigned long long file_size = 0;
     unsigned long long part_begin = 0;
     unsigned long long part_end = 0;
     unsigned long long part_size = 0;
@@ -113,6 +114,12 @@ PyObject* yenc_decode(PyObject* self, PyObject* Py_bytesarray_obj) {
         retval = NULL;
         goto finish;
     }
+
+    // Get the size of the reconstructed file
+    start_loc = my_memstr(start_loc, end_loc - start_loc, "size=", 1);
+    if (start_loc) {
+        file_size = atoll(start_loc);
+    }
 
     // Find start of the filename
     start_loc = my_memstr(start_loc, end_loc - start_loc, " name=", 1);
@@ -215,7 +222,7 @@ PyObject* yenc_decode(PyObject* self, PyObject* Py_bytesarray_obj) {
     }
 
     // Build output
-    retval = Py_BuildValue("(S, K, K, N)", Py_output_filename, part_begin, part_size, Py_output_crc);
+    retval = Py_BuildValue("(S, K, K, K, N)", Py_output_filename, file_size, part_begin, part_size, Py_output_crc);
 
 finish:
     Py_XDECREF(Py_output_filename);

tests/test_decoder.py

@@ -13,8 +13,9 @@ def test_regular():
 
 def test_partial():
     data_plain = read_plain_yenc_file("test_partial.yenc")
-    decoded_data, filename, begin, size, crc_correct = sabctools_yenc_wrapper(data_plain)
+    decoded_data, filename, filesize, begin, size, crc_correct = sabctools_yenc_wrapper(data_plain)
     assert filename == "90E2Sdvsmds0801dvsmds90E.part06.rar"
+    assert filesize == 49152000
     assert begin == 15360000
     assert size == 384000
     assert crc_correct is None
@@ -67,7 +68,7 @@ def test_ref_counts():
     """Note that sys.getrefcount itself adds another reference!"""
     # Test regular case
     data_plain = read_plain_yenc_file("test_regular.yenc")
-    data_out, filename, begin, end, crc_correct = sabctools_yenc_wrapper(data_plain)
+    data_out, filename, filesize, begin, end, crc_correct = sabctools_yenc_wrapper(data_plain)
     # data_plain and data_out point to the same data!
     assert sys.getrefcount(data_plain) == 3
     assert sys.getrefcount(data_out) == 3

tests/testsupport.py

@@ -66,8 +66,8 @@ def read_pickle(filename):
 
 
 def sabctools_yenc_wrapper(data: bytearray) -> Tuple[bytearray, str, int, int, Optional[int]]:
-    filename, begin, size, crc_correct = sabctools.yenc_decode(data)
-    return data, correct_unknown_encoding(filename), begin, size, crc_correct
+    filename, filesize, begin, size, crc_correct = sabctools.yenc_decode(data)
+    return data, correct_unknown_encoding(filename), filesize, begin, size, crc_correct
 
 
 def python_yenc(data_plain):
@@ -112,7 +112,7 @@ def python_yenc(data_plain):
             size = end - begin + 1
             begin -= 1
 
-    return decoded_data, ybegin["name"], begin, size, binascii.crc32(decoded_data)
+    return decoded_data, ybegin["name"], int(ybegin["size"]), begin, size, binascii.crc32(decoded_data)
 
 
 def parse_yenc_data(data):