Skip to content

Commit 13f4fad

Browse files
committed
Refactor zipfile._strip_extra to use higher level abstractions for extras instead of a heavy-state loop.
1 parent 32df540 commit 13f4fad

File tree

1 file changed

+37
-20
lines changed

1 file changed

+37
-20
lines changed

Lib/zipfile/__init__.py

+37-20
Original file line numberDiff line numberDiff line change
@@ -188,26 +188,43 @@ class LargeZipFile(Exception):
188188

189189
_DD_SIGNATURE = 0x08074b50
190190

191-
_EXTRA_FIELD_STRUCT = struct.Struct('<HH')
192-
193-
def _strip_extra(extra, xids):
194-
# Remove Extra Fields with specified IDs.
195-
unpack = _EXTRA_FIELD_STRUCT.unpack
196-
modified = False
197-
buffer = []
198-
start = i = 0
199-
while i + 4 <= len(extra):
200-
xid, xlen = unpack(extra[i : i + 4])
201-
j = i + 4 + xlen
202-
if xid in xids:
203-
if i != start:
204-
buffer.append(extra[start : i])
205-
start = j
206-
modified = True
207-
i = j
208-
if not modified:
209-
return extra
210-
return b''.join(buffer)
191+
192+
class _Extra(bytes):
193+
FIELD_STRUCT = struct.Struct('<HH')
194+
195+
def __new__(cls, val, id=None):
196+
return super().__new__(cls, val)
197+
198+
def __init__(self, val, id=None):
199+
self.id = id
200+
201+
@classmethod
202+
def read_one(cls, raw):
203+
try:
204+
xid, xlen = cls.FIELD_STRUCT.unpack(raw[:4])
205+
except struct.error:
206+
xid = None
207+
xlen = 0
208+
return cls(raw[:4+xlen], xid), raw[4+xlen:]
209+
210+
@classmethod
211+
def split(cls, data):
212+
while data:
213+
extra, data = _Extra.read_one(data)
214+
yield extra
215+
216+
@classmethod
217+
def strip(cls, data, xids):
218+
"""Remove Extra fields with specified IDs."""
219+
return b''.join(
220+
ex
221+
for ex in cls.split(data)
222+
if ex.id not in xids
223+
)
224+
225+
226+
_strip_extra = _Extra.strip
227+
211228

212229
def _check_zipfile(fp):
213230
try:

0 commit comments

Comments
 (0)