gh-152140: Rework and optimize _Extra class in zipfile#152141
gh-152140: Rework and optimize _Extra class in zipfile#152141danny0838 wants to merge 8 commits into
_Extra class in zipfile#152141Conversation
eendebakpt
left a comment
There was a problem hiding this comment.
The PR indeed creates cleaner code. But this is not purely refactoring, there are several new things:
- Early exit for empty extra data: does that occur? If so, how often
- _strip_extra_fields now returns a bytarray instead of bytes
- The parsing of the data is implemented different now (e.g. we need an additional
keep remaining trailing bytes)
Can you make benchmarks to show the new code is indeed faster (or at least no regression)?
It depends. Other ZIP tools may add Linux permission or extended date, while Python
Yes, this doesn't affect the current single use case. And it would be better for future potential caller to further process the returned value.
This is NOT a behavior change, as responsed in another comment.
I'll try to. Is there a need to include it in the test suite? Or just provide a standalone code sinipper for testing? |
|
Here's the benchmark: import timeit
import struct
import random
random.seed(42)
class _Extra(bytes):
FIELD_STRUCT = struct.Struct('<HH')
def __new__(cls, val, id=None):
return super().__new__(cls, val)
def __init__(self, val, id=None):
self.id = id
@classmethod
def read_one(cls, raw):
try:
xid, xlen = cls.FIELD_STRUCT.unpack(raw[:4])
except struct.error:
xid = None
xlen = 0
return cls(raw[:4+xlen], xid), raw[4+xlen:]
@classmethod
def split(cls, data):
# use memoryview for zero-copy slices
rest = memoryview(data)
while rest:
extra, rest = _Extra.read_one(rest)
yield extra
@classmethod
def strip(cls, data, xids):
"""Remove Extra fields with specified IDs."""
return b''.join(
ex
for ex in cls.split(data)
if ex.id not in xids
)
def strip_extra_fields(data, field_ids):
"""Remove Extra fields with specified IDs and return a bytearray.
data should be bytes or bytearray.
"""
result = bytearray()
# early return for empty extra data
if not data:
return result
# use memoryview for zero-copy slices
mv = memoryview(data)
mv_len = len(mv)
pos = 0
while pos + 4 <= mv_len:
xid, xlen = struct.unpack_from('<HH', mv, pos)
if pos + 4 + xlen > mv_len:
break
if xid not in field_ids:
result.extend(mv[pos:pos + 4 + xlen])
pos += 4 + xlen
# keep remaining trailing bytes (e.g. truncated or malformed data)
if pos < mv_len:
result.extend(mv[pos:])
return result
def strip_extra_fields_no_early_return(data, field_ids):
"""Remove Extra fields with specified IDs and return a bytearray.
data should be bytes or bytearray.
"""
result = bytearray()
# use memoryview for zero-copy slices
mv = memoryview(data)
mv_len = len(mv)
pos = 0
while pos + 4 <= mv_len:
xid, xlen = struct.unpack_from('<HH', mv, pos)
if pos + 4 + xlen > mv_len:
break
if xid not in field_ids:
result.extend(mv[pos:pos + 4 + xlen])
pos += 4 + xlen
# keep remaining trailing bytes (e.g. truncated or malformed data)
if pos < mv_len:
result.extend(mv[pos:])
return result
def strip_extra_fields_no_mv(data, field_ids):
"""Remove Extra fields with specified IDs and return a bytearray.
data should be bytes or bytearray.
"""
result = bytearray()
# early return for empty extra data
if not data:
return result
mv = data
mv_len = len(mv)
pos = 0
while pos + 4 <= mv_len:
xid, xlen = struct.unpack_from('<HH', mv, pos)
if pos + 4 + xlen > mv_len:
break
if xid not in field_ids:
result.extend(mv[pos:pos + 4 + xlen])
pos += 4 + xlen
# keep remaining trailing bytes (e.g. truncated or malformed data)
if pos < mv_len:
result.extend(mv[pos:])
return result
def strip_extra_fields_no_mv_no_early_return(data, field_ids):
"""Remove Extra fields with specified IDs and return a bytearray.
data should be bytes or bytearray.
"""
result = bytearray()
mv = data
mv_len = len(mv)
pos = 0
while pos + 4 <= mv_len:
xid, xlen = struct.unpack_from('<HH', mv, pos)
if pos + 4 + xlen > mv_len:
break
if xid not in field_ids:
result.extend(mv[pos:pos + 4 + xlen])
pos += 4 + xlen
# keep remaining trailing bytes (e.g. truncated or malformed data)
if pos < mv_len:
result.extend(mv[pos:])
return result
def strip_extra_fields_try(data, field_ids):
"""Remove Extra fields with specified IDs and return a bytearray.
data should be bytes or bytearray.
"""
result = bytearray()
# early return for empty extra data
if not data:
return result
# use memoryview for zero-copy slices
data_len = len(data)
pos = 0
while pos < data_len:
try:
xid, xlen = struct.unpack_from('<HH', data, pos)
except struct.error:
xid, xlen = None, 0
if xid not in field_ids:
result.extend(data[pos:pos + 4 + xlen])
pos += 4 + xlen
return result
class _Extra2:
@classmethod
def iter(cls, data):
"""Iter through all subfields."""
# early return for empty extra data
if not data:
return
pos, data_len = 0, len(data)
while pos < data_len:
try:
xid, xlen = struct.unpack_from('<HH', data, pos)
except struct.error:
xid, xlen = None, 0
yield data[pos:pos + 4 + xlen], xid, xlen
pos += 4 + xlen
@classmethod
def strip(cls, data, xids):
"""Remove Extra fields with specified IDs and return a bytearray."""
result = bytearray()
for xfield, xid, _ in cls.iter(data):
if xid not in xids:
result.extend(xfield)
return result
class _Extra3:
@classmethod
def iter(cls, data):
"""Iter through all subfields."""
# early return for empty extra data
if not data:
return
pos, data_len = 0, len(data)
while pos < data_len:
try:
xid, xlen = struct.unpack_from('<HH', data, pos)
except struct.error:
xid, xlen = None, 0
yield data[pos:pos + 4 + xlen], xid, xlen
pos += 4 + xlen
@classmethod
def strip(cls, data, xids):
"""Remove Extra fields with specified IDs and return a bytearray."""
result = b''
for xfield, xid, _ in cls.iter(data):
if xid not in xids:
result += xfield
return result
class _Extra4:
@classmethod
def iter(cls, data):
"""Iter through all subfields."""
# early return for empty extra data
if not data:
return
pos, data_len = 0, len(data)
while pos < data_len:
try:
xid, xlen = struct.unpack_from('<HH', data, pos)
except struct.error:
xid, xlen = None, 0
yield data[pos:pos + 4 + xlen], xid, xlen
pos += 4 + xlen
@classmethod
def strip(cls, data, xids):
"""Remove Extra fields with specified IDs and return a bytearray."""
return b''.join(
xfield
for xfield, xid, _ in cls.iter(data)
if xid not in xids
)
class _Extra5:
@classmethod
def iter(cls, data):
"""Iter through all subfields."""
# early return for empty extra data
if not data:
return
pos, data_len = 0, len(data)
while pos < data_len:
try:
xid, xlen = struct.unpack_from('<HH', data, pos)
except struct.error:
xid, xlen = None, 0
yield data[pos:pos + 4 + xlen], xid
pos += 4 + xlen
@classmethod
def strip(cls, data, xids):
"""Remove Extra fields with specified IDs and return a bytearray."""
result = b''
for xfield, xid in cls.iter(data):
if xid not in xids:
result += xfield
return result
class _Extra6:
FIELD_STRUCT = struct.Struct('<HH')
@classmethod
def iter(cls, data):
"""Iter through all subfields."""
# early return for empty extra data
if not data:
return
pos, data_len = 0, len(data)
while pos < data_len:
try:
xid, xlen = cls.FIELD_STRUCT.unpack_from(data, pos)
except struct.error:
xid, xlen = None, 0
yield data[pos:pos + 4 + xlen], xid, xlen
pos += 4 + xlen
@classmethod
def strip(cls, data, xids):
"""Remove Extra fields with specified IDs."""
result = b''
for xfield, xid, _ in cls.iter(data):
if xid not in xids:
result += xfield
return result
s = struct.Struct("<HH")
valid_extra_data = (
s.pack(0x0001, 24) + b"X" * 24 + # ZIP64 Extended Information
s.pack(0x5455, 13) + b"\x07" + struct.pack("<LLL", 1767225600, 1767225600, 1767225600) + # Extended Timestamp
s.pack(0x000a, 32) + b"\x00"*4 + struct.pack("<QQQ", 132537600000000000, 132537600000000000, 132537600000000000) + # NTFS Info
s.pack(0x0001, 16) + b"Y" * 16 + # ZIP64
s.pack(0x7875, 11) + b"\x01\x04\xe8\x03\x00\x00\x04\xe8\x03\x00\x00" # Unix UID/GID
)
empty_extra_data = b""
dataset = [valid_extra_data if i % 2 == 0 else empty_extra_data for i in range(200)]
random.shuffle(dataset)
def run_case0():
for data in dataset:
_Extra.strip(data, (1,))
def run_case1():
for data in dataset:
strip_extra_fields(data, (1,))
def run_case2():
for data in dataset:
strip_extra_fields_no_early_return(data, (1,))
def run_case3():
for data in dataset:
strip_extra_fields_no_mv(data, (1,))
def run_case4():
for data in dataset:
strip_extra_fields_no_mv_no_early_return(data, (1,))
def run_case5():
for data in dataset:
strip_extra_fields_try(data, (1,))
def run_case6():
for data in dataset:
_Extra2.strip(data, (1,))
def run_case7():
for data in dataset:
_Extra3.strip(data, (1,))
def run_case8():
for data in dataset:
_Extra4.strip(data, (1,))
def run_case9():
for data in dataset:
_Extra5.strip(data, (1,))
def run_case10():
for data in dataset:
_Extra6.strip(data, (1,))
NUMBER = 10000
print(f"=== Benchmark Results ({NUMBER * len(dataset):,} total invocations) ===")
print(f"Dataset Profile: 50% Valid Extra Data, 50% Empty Data")
i = 0
while True:
try:
time = timeit.timeit(stmt=f'run_case{i}()', globals=globals(), number=NUMBER)
except NameError:
break
if i == 0:
time0 = time
print(f'case{i}: {time:.4f} seconds ({time0 / time:.2f}x)')
i += 1
import tracemalloc
def profile_memory(func):
tracemalloc.start()
for _ in range(50):
func()
current, peak = tracemalloc.get_traced_memory()
tracemalloc.stop()
return peak
print(f"\n=== Memory Peak Results ===")
i = 0
while True:
try:
peak = profile_memory(globals()[f'run_case{i}'])
except KeyError:
break
if i == 0:
peak0 = peak
print(
f'Peak Memory for case{i}: {peak:,} bytes '
f'(saved {((peak0 - peak) / peak0) * 100:.1f}%)'
)
i += 1And the results: The benchmark shows that the new implementation improves performance, the early return for empty data is helpful (assume that 50% is empty), removing |
_Extra class with ZipFile._strip_extra_fields()_Extra class with _strip_extra_fields() in zipfile
_Extra class with _strip_extra_fields() in zipfile_Extra class in zipfile
Introduce a helper `.update()` that simplifies the common operation of fields updating of extra data bytes, by automatically stripping existing subfields with identical IDs and inserts multiple fields in one step.
Make `_Extra` a normal class rather than a subclass of `bytes`, avoiding the performance penalty of copying the passed bytes upon each instantiation. Introduce `.iter()` in place of `.strip()` to yield each subfield info as a tuple. This provides better semantics and is more performant than subclassing `bytes`. The method name change also enables feature detection for these behavior changes. Stop using memoryview internally because each slice creates an overhead of around 184 bytes, making it inefficient for small bytes operations (each extra field is usually only several to tens of bytes). Remove `.read_one()` since it is now unused and obsolete.
|
Rework the There is currently no significant performance impact since the only caller is involved when writing a really large file, making the overhead of stripping extras relatively small. However, there will be an impact if gh-152487 and gh-152704 are merged, you can see gh-152140-2 and checkout a3a54a1 and 409a733 for pre-patch and post-patch benchmark respectively. The benchmark code: import timeit
import os
import io
import struct
from zipfile import ZipFile, ZipInfo, _Extra
TESTFN = "bench_temp_file.txt"
with open(TESTFN, "w") as f:
f.write("bench_data")
s = struct.Struct("<HH")
existing_extra = s.pack(0x5455, 5) + b"A"*5 + s.pack(0x000a, 32) + b"B"*32
existing_extra_z64 = s.pack(0x0001, 24) + b"Z"*24 + existing_extra
def bench_force_zip64_with_existing_extra():
bio = io.BytesIO()
with ZipFile(bio, mode="w") as zf:
for i in range(500):
zinfo = ZipInfo(f"file_{i}.txt")
# assume 1% contains existing zip64 data
zinfo.extra = existing_extra_z64 if i % 100 == 0 else existing_extra
with zf.open(zinfo, mode="w", force_zip64=True) as f:
f.write(b"x")
def bench_force_zip64_with_ext_timestamps():
bio = io.BytesIO()
with ZipFile(bio, mode="w", with_ext_timestamps=True) as zf:
for i in range(500):
zinfo = ZipInfo.from_file(TESTFN, f"file_{i}.txt")
with zf.open(zinfo, mode="w", force_zip64=True) as f:
f.write(b"foo")
f.write(b"bar")
def main():
patched = '(PATCHED)' if hasattr(_Extra, 'iter') else '(non-PATCHED)'
print("=" * 65)
print(f"ZIPFILE OPEN API PERFORMANCE BENCHMARK {patched}")
print("=" * 65)
loops = 500
total_files = loops * 500
time = timeit.timeit(bench_force_zip64_with_existing_extra, number=loops)
print(f"[Case] ZipFile.open(force_zip64=True) with existing extra (1% zip64)")
print(f" Loops: {loops} ({total_files} total files written)")
print(f" Total Time: {time:.4f} seconds")
print(f" Avg Per Loop: {time / loops * 1e3:.2f} ms")
print(f" Avg Per File: {time / total_files * 1e6:.2f} μs")
print("-" * 65)
loops = 200
total_files = loops * 500
time = timeit.timeit(bench_force_zip64_with_ext_timestamps, number=loops)
print(f"[Case] ZipFile.open(force_zip64=True) with ZipInfo.from_file(with_ext_timestamps=True)")
print(f" Loops: {loops} ({total_files} total files written)")
print(f" Total Time: {time:.4f} seconds")
print(f" Avg Per Loop: {time / loops * 1e6:.2f} μs")
print(f" Avg Per File: {time / total_files * 1e6:.2f} μs")
print("-" * 65)
if os.path.exists(TESTFN):
os.remove(TESTFN)
if __name__ == "__main__":
main()The results tested on cmd and powershell in Window 10 are: It shows grossly 1.06~1.21x faster for writing massive small files with |
Use `unpack_from` with offset to prevent bytes slicing.
The
_Extraclass was over-engineered. Its only active usage was itsstrip()method, called byZipFile._write_end_record()to provide the stripping logic for ZIP64 fields. The context-dependent nature of extra fields also made it difficult to be reused by_decodeExtra()or other methods efficiently. Additionally, itssplit()method called_Extradirectly rather than utilizingcls, which was a suboptimal pattern that hindered extensibility.Remove the
_Extraclass entirely and reimplement it as a private static method_strip_extra_fields()that processes a bytearray insideZipFile, positioned directly beneath its caller. This eliminates dead and suboptimal code, achieves clean encapsulation and code locality, and improves performance by avoiding temporary class allocations.Additionally, bind this method as a class property in the tests to simplify the code.
zipfilemodule #152140