Skip to content

Commit 2ff724f

Browse files
danny0838danny0838
committed
gh-152140: Replace _Extra class with _strip_extra_fields() in zipifle
The `_Extra` class was over-engineered. Its only active usage was its `strip()` method, called by `ZipFile._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, its `split()` method called `_Extra` directly rather than utilizing `cls`, which was a suboptimal pattern that hindered extensibility. Remove the `_Extra` class entirely and reimplement it as a module function `_strip_extra_fields()` that processes a bytearray. This eliminates dead and suboptimal code and improves performance by avoiding temporary class allocations.
1 parent 65585ca commit 2ff724f

3 files changed

Lines changed: 64 additions & 69 deletions

File tree

‎Lib/test/test_zipfile/test_core.py‎

Lines changed: 39 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -5840,59 +5840,69 @@ class StripExtraTests(unittest.TestCase):
58405840

58415841
ZIP64_EXTRA = 1
58425842

5843+
strip_extra = staticmethod(zipfile._strip_extra_fields)
5844+
58435845
def test_no_data(self):
58445846
s = struct.Struct("<HH")
58455847
a = s.pack(self.ZIP64_EXTRA, 0)
58465848
b = s.pack(2, 0)
58475849
c = s.pack(3, 0)
58485850

5849-
self.assertEqual(b'', zipfile._Extra.strip(a, (self.ZIP64_EXTRA,)))
5850-
self.assertEqual(b, zipfile._Extra.strip(b, (self.ZIP64_EXTRA,)))
5851-
self.assertEqual(
5852-
b+b"z", zipfile._Extra.strip(b+b"z", (self.ZIP64_EXTRA,)))
5851+
self.assertEqual(b'', self.strip_extra(a, (self.ZIP64_EXTRA,)))
5852+
self.assertEqual(b, self.strip_extra(b, (self.ZIP64_EXTRA,)))
5853+
self.assertEqual(b+b"z", self.strip_extra(b+b"z", (self.ZIP64_EXTRA,)))
58535854

5854-
self.assertEqual(b+c, zipfile._Extra.strip(a+b+c, (self.ZIP64_EXTRA,)))
5855-
self.assertEqual(b+c, zipfile._Extra.strip(b+a+c, (self.ZIP64_EXTRA,)))
5856-
self.assertEqual(b+c, zipfile._Extra.strip(b+c+a, (self.ZIP64_EXTRA,)))
5855+
self.assertEqual(b+c, self.strip_extra(a+b+c, (self.ZIP64_EXTRA,)))
5856+
self.assertEqual(b+c, self.strip_extra(b+a+c, (self.ZIP64_EXTRA,)))
5857+
self.assertEqual(b+c, self.strip_extra(b+c+a, (self.ZIP64_EXTRA,)))
58575858

58585859
def test_with_data(self):
58595860
s = struct.Struct("<HH")
58605861
a = s.pack(self.ZIP64_EXTRA, 1) + b"a"
58615862
b = s.pack(2, 2) + b"bb"
58625863
c = s.pack(3, 3) + b"ccc"
58635864

5864-
self.assertEqual(b"", zipfile._Extra.strip(a, (self.ZIP64_EXTRA,)))
5865-
self.assertEqual(b, zipfile._Extra.strip(b, (self.ZIP64_EXTRA,)))
5866-
self.assertEqual(
5867-
b+b"z", zipfile._Extra.strip(b+b"z", (self.ZIP64_EXTRA,)))
5865+
self.assertEqual(b"", self.strip_extra(a, (self.ZIP64_EXTRA,)))
5866+
self.assertEqual(b, self.strip_extra(b, (self.ZIP64_EXTRA,)))
5867+
self.assertEqual(b+b"z", self.strip_extra(b+b"z", (self.ZIP64_EXTRA,)))
58685868

5869-
self.assertEqual(b+c, zipfile._Extra.strip(a+b+c, (self.ZIP64_EXTRA,)))
5870-
self.assertEqual(b+c, zipfile._Extra.strip(b+a+c, (self.ZIP64_EXTRA,)))
5871-
self.assertEqual(b+c, zipfile._Extra.strip(b+c+a, (self.ZIP64_EXTRA,)))
5869+
self.assertEqual(b+c, self.strip_extra(a+b+c, (self.ZIP64_EXTRA,)))
5870+
self.assertEqual(b+c, self.strip_extra(b+a+c, (self.ZIP64_EXTRA,)))
5871+
self.assertEqual(b+c, self.strip_extra(b+c+a, (self.ZIP64_EXTRA,)))
58725872

58735873
def test_multiples(self):
58745874
s = struct.Struct("<HH")
58755875
a = s.pack(self.ZIP64_EXTRA, 1) + b"a"
58765876
b = s.pack(2, 2) + b"bb"
58775877

5878-
self.assertEqual(b"", zipfile._Extra.strip(a+a, (self.ZIP64_EXTRA,)))
5879-
self.assertEqual(b"", zipfile._Extra.strip(a+a+a, (self.ZIP64_EXTRA,)))
5880-
self.assertEqual(
5881-
b"z", zipfile._Extra.strip(a+a+b"z", (self.ZIP64_EXTRA,)))
5882-
self.assertEqual(
5883-
b+b"z", zipfile._Extra.strip(a+a+b+b"z", (self.ZIP64_EXTRA,)))
5878+
self.assertEqual(b"", self.strip_extra(a+a, (self.ZIP64_EXTRA,)))
5879+
self.assertEqual(b"", self.strip_extra(a+a+a, (self.ZIP64_EXTRA,)))
5880+
self.assertEqual(b"z", self.strip_extra(a+a+b"z", (self.ZIP64_EXTRA,)))
5881+
self.assertEqual(b+b"z", self.strip_extra(a+a+b+b"z", (self.ZIP64_EXTRA,)))
5882+
5883+
self.assertEqual(b, self.strip_extra(a+a+b, (self.ZIP64_EXTRA,)))
5884+
self.assertEqual(b, self.strip_extra(a+b+a, (self.ZIP64_EXTRA,)))
5885+
self.assertEqual(b, self.strip_extra(b+a+a, (self.ZIP64_EXTRA,)))
58845886

5885-
self.assertEqual(b, zipfile._Extra.strip(a+a+b, (self.ZIP64_EXTRA,)))
5886-
self.assertEqual(b, zipfile._Extra.strip(a+b+a, (self.ZIP64_EXTRA,)))
5887-
self.assertEqual(b, zipfile._Extra.strip(b+a+a, (self.ZIP64_EXTRA,)))
5887+
def test_truncated_data_for_stripped_id(self):
5888+
s = struct.Struct("<HH")
5889+
a = s.pack(self.ZIP64_EXTRA, 3) + b"a"
5890+
b = s.pack(2, 2) + b"bb"
5891+
5892+
self.assertEqual(b, self.strip_extra(b+a, (self.ZIP64_EXTRA,)))
5893+
5894+
def test_truncated_data_for_nonstripped_id(self):
5895+
s = struct.Struct("<HH")
5896+
a = s.pack(self.ZIP64_EXTRA, 1) + b"a"
5897+
b = s.pack(2, 4) + b"bb"
5898+
5899+
self.assertEqual(b, self.strip_extra(a+b, (self.ZIP64_EXTRA,)))
58885900

58895901
def test_too_short(self):
5890-
self.assertEqual(b"", zipfile._Extra.strip(b"", (self.ZIP64_EXTRA,)))
5891-
self.assertEqual(b"z", zipfile._Extra.strip(b"z", (self.ZIP64_EXTRA,)))
5892-
self.assertEqual(
5893-
b"zz", zipfile._Extra.strip(b"zz", (self.ZIP64_EXTRA,)))
5894-
self.assertEqual(
5895-
b"zzz", zipfile._Extra.strip(b"zzz", (self.ZIP64_EXTRA,)))
5902+
self.assertEqual(b"", self.strip_extra(b"", (self.ZIP64_EXTRA,)))
5903+
self.assertEqual(b"z", self.strip_extra(b"z", (self.ZIP64_EXTRA,)))
5904+
self.assertEqual(b"zz", self.strip_extra(b"zz", (self.ZIP64_EXTRA,)))
5905+
self.assertEqual(b"zzz", self.strip_extra(b"zzz", (self.ZIP64_EXTRA,)))
58965906

58975907

58985908
class StatIO(_pyio.BytesIO):

‎Lib/zipfile/__init__.py‎

Lines changed: 24 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -194,43 +194,6 @@ class LargeZipFile(Exception):
194194

195195
_DD_SIGNATURE = 0x08074b50
196196

197-
198-
class _Extra(bytes):
199-
FIELD_STRUCT = struct.Struct('<HH')
200-
201-
def __new__(cls, val, id=None):
202-
return super().__new__(cls, val)
203-
204-
def __init__(self, val, id=None):
205-
self.id = id
206-
207-
@classmethod
208-
def read_one(cls, raw):
209-
try:
210-
xid, xlen = cls.FIELD_STRUCT.unpack(raw[:4])
211-
except struct.error:
212-
xid = None
213-
xlen = 0
214-
return cls(raw[:4+xlen], xid), raw[4+xlen:]
215-
216-
@classmethod
217-
def split(cls, data):
218-
# use memoryview for zero-copy slices
219-
rest = memoryview(data)
220-
while rest:
221-
extra, rest = _Extra.read_one(rest)
222-
yield extra
223-
224-
@classmethod
225-
def strip(cls, data, xids):
226-
"""Remove Extra fields with specified IDs."""
227-
return b''.join(
228-
ex
229-
for ex in cls.split(data)
230-
if ex.id not in xids
231-
)
232-
233-
234197
def _check_zipfile(fp):
235198
try:
236199
endrec = _EndRecData(fp)
@@ -427,6 +390,27 @@ def _read_local_file_header(fp):
427390
raise BadZipFile("Bad magic number for file header")
428391
return fheader
429392

393+
def _strip_extra_fields(data, field_ids):
394+
"""Remove Extra fields with specified IDs and return a bytearray."""
395+
result = bytearray()
396+
397+
# early return for empty extra data
398+
if not data:
399+
return result
400+
401+
data_len = len(data)
402+
pos = 0
403+
while pos < data_len:
404+
try:
405+
xid, xlen = struct.unpack_from('<HH', data, pos)
406+
except struct.error:
407+
xid, xlen = None, 0
408+
if xid not in field_ids:
409+
result.extend(data[pos:pos + 4 + xlen])
410+
pos += 4 + xlen
411+
412+
return result
413+
430414

431415
class ZipInfo:
432416
"""Class with attributes describing each file in the ZIP archive."""
@@ -2665,11 +2649,11 @@ def _write_end_record(self):
26652649
extra_data = zinfo.extra
26662650
min_version = 0
26672651
if extra:
2668-
# Append a ZIP64 field to the extra's
2669-
extra_data = _Extra.strip(extra_data, (1,))
2652+
# Prepend a ZIP64 field to the extra's
26702653
extra_data = struct.pack(
26712654
'<HH' + 'Q'*len(extra),
2672-
1, 8*len(extra), *extra) + extra_data
2655+
1, 8*len(extra), *extra
2656+
) + _strip_extra_fields(extra_data, (1,))
26732657

26742658
min_version = ZIP64_VERSION
26752659

‎Misc/NEWS.d/next/Core_and_Builtins/2026-06-25-00-41-52.gh-issue-152140.u0phBe.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Replace :class:`!_Extra` with :func:`!_strip_extra_fields` in the :mod:`zipfile` module.

0 commit comments

Comments
 (0)