Skip to content

Commit e6267ad

Browse files
danny0838danny0838
committed
gh-152140: Replace _Extra class with ZipFile._strip_extra_fields()
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 private static method `_strip_extra_fields()` that processes a bytearray inside `ZipFile`, 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.
1 parent 5049eb1 commit e6267ad

2 files changed

Lines changed: 60 additions & 69 deletions

File tree

‎Lib/test/test_zipfile/test_core.py‎

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

58415841
ZIP64_EXTRA = 1
58425842

5843+
@classmethod
5844+
def setUpClass(cls):
5845+
cls.strip_extra = staticmethod(zipfile.ZipFile._strip_extra_fields)
5846+
58435847
def test_no_data(self):
58445848
s = struct.Struct("<HH")
58455849
a = s.pack(self.ZIP64_EXTRA, 0)
58465850
b = s.pack(2, 0)
58475851
c = s.pack(3, 0)
58485852

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,)))
5853+
self.assertEqual(b'', self.strip_extra(a, (self.ZIP64_EXTRA,)))
5854+
self.assertEqual(b, self.strip_extra(b, (self.ZIP64_EXTRA,)))
5855+
self.assertEqual(b+b"z", self.strip_extra(b+b"z", (self.ZIP64_EXTRA,)))
58535856

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,)))
5857+
self.assertEqual(b+c, self.strip_extra(a+b+c, (self.ZIP64_EXTRA,)))
5858+
self.assertEqual(b+c, self.strip_extra(b+a+c, (self.ZIP64_EXTRA,)))
5859+
self.assertEqual(b+c, self.strip_extra(b+c+a, (self.ZIP64_EXTRA,)))
58575860

58585861
def test_with_data(self):
58595862
s = struct.Struct("<HH")
58605863
a = s.pack(self.ZIP64_EXTRA, 1) + b"a"
58615864
b = s.pack(2, 2) + b"bb"
58625865
c = s.pack(3, 3) + b"ccc"
58635866

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,)))
5867+
self.assertEqual(b"", self.strip_extra(a, (self.ZIP64_EXTRA,)))
5868+
self.assertEqual(b, self.strip_extra(b, (self.ZIP64_EXTRA,)))
5869+
self.assertEqual(b+b"z", self.strip_extra(b+b"z", (self.ZIP64_EXTRA,)))
58685870

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,)))
5871+
self.assertEqual(b+c, self.strip_extra(a+b+c, (self.ZIP64_EXTRA,)))
5872+
self.assertEqual(b+c, self.strip_extra(b+a+c, (self.ZIP64_EXTRA,)))
5873+
self.assertEqual(b+c, self.strip_extra(b+c+a, (self.ZIP64_EXTRA,)))
58725874

58735875
def test_multiples(self):
58745876
s = struct.Struct("<HH")
58755877
a = s.pack(self.ZIP64_EXTRA, 1) + b"a"
58765878
b = s.pack(2, 2) + b"bb"
58775879

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,)))
5880+
self.assertEqual(b"", self.strip_extra(a+a, (self.ZIP64_EXTRA,)))
5881+
self.assertEqual(b"", self.strip_extra(a+a+a, (self.ZIP64_EXTRA,)))
5882+
self.assertEqual(b"z", self.strip_extra(a+a+b"z", (self.ZIP64_EXTRA,)))
5883+
self.assertEqual(b+b"z", self.strip_extra(a+a+b+b"z", (self.ZIP64_EXTRA,)))
58845884

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,)))
5885+
self.assertEqual(b, self.strip_extra(a+a+b, (self.ZIP64_EXTRA,)))
5886+
self.assertEqual(b, self.strip_extra(a+b+a, (self.ZIP64_EXTRA,)))
5887+
self.assertEqual(b, self.strip_extra(b+a+a, (self.ZIP64_EXTRA,)))
58885888

58895889
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,)))
5890+
self.assertEqual(b"", self.strip_extra(b"", (self.ZIP64_EXTRA,)))
5891+
self.assertEqual(b"z", self.strip_extra(b"z", (self.ZIP64_EXTRA,)))
5892+
self.assertEqual(b"zz", self.strip_extra(b"zz", (self.ZIP64_EXTRA,)))
5893+
self.assertEqual(b"zzz", self.strip_extra(b"zzz", (self.ZIP64_EXTRA,)))
58965894

58975895

58985896
class StatIO(_pyio.BytesIO):

‎Lib/zipfile/__init__.py‎

Lines changed: 33 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)
@@ -2670,11 +2633,11 @@ def _write_end_record(self):
26702633
extra_data = zinfo.extra
26712634
min_version = 0
26722635
if extra:
2673-
# Append a ZIP64 field to the extra's
2674-
extra_data = _Extra.strip(extra_data, (1,))
2636+
# Prepend a ZIP64 field to the extra's
26752637
extra_data = struct.pack(
26762638
'<HH' + 'Q'*len(extra),
2677-
1, 8*len(extra), *extra) + extra_data
2639+
1, 8*len(extra), *extra
2640+
) + self._strip_extra_fields(extra_data, (1,))
26782641

26792642
min_version = ZIP64_VERSION
26802643

@@ -2741,6 +2704,36 @@ def _write_end_record(self):
27412704
self.fp.truncate()
27422705
self.fp.flush()
27432706

2707+
@staticmethod
2708+
def _strip_extra_fields(data, field_ids):
2709+
"""Remove Extra fields with specified IDs and return a bytearray.
2710+
2711+
data should be bytes or bytearray.
2712+
"""
2713+
result = bytearray()
2714+
2715+
# early return for empty extra data
2716+
if not data:
2717+
return result
2718+
2719+
# use memoryview for zero-copy slices
2720+
mv = memoryview(data)
2721+
mv_len = len(mv)
2722+
pos = 0
2723+
while pos + 4 <= mv_len:
2724+
xid, xlen = struct.unpack_from('<HH', mv, pos)
2725+
if pos + 4 + xlen > mv_len:
2726+
break
2727+
if xid not in field_ids:
2728+
result.extend(mv[pos:pos + 4 + xlen])
2729+
pos += 4 + xlen
2730+
2731+
# keep remaining trailing bytes (e.g. truncated or malformed data)
2732+
if pos < mv_len:
2733+
result.extend(mv[pos:])
2734+
2735+
return result
2736+
27442737
def _fpclose(self, fp):
27452738
assert self._fileRefCnt > 0
27462739
self._fileRefCnt -= 1

0 commit comments

Comments
 (0)